def choose_max(self, axis, target = None):
        """
        Sets the argmax along axis to 1 and rest to zero.
        """

        m, n = self.shape

        assert axis == 0, 'Axis = 1 not implemented.'
        if not target:
          target = self

        err_code =  _cudamat.choose_max_by_axis(self.p_mat, target.p_mat, ct.c_int(axis))
        if err_code:
            raise generate_exception(err_code)

        return target

        Sets the argmax along axis to 1 and rest to zero.
        
cudamat/cudamat.py