fix apis for d class (#2151)

Author: lvyufeng

mindnlp/core/_prims/ascend.py

@@ -332,3 +332,46 @@ def grid_sampler_3d(input, grid, mode, padding_mode, align_corners):
 
 __all__.append('grid_sampler_2d')
 __all__.append('grid_sampler_3d')
+
+def rms_norm(x, gamma, epsilon):
+    return pyboost_inner_prim.rms_norm_impl(x, gamma, epsilon)[0]
+
+__all__.append('rms_norm')
+
+_complex = ops.Complex().set_device('Ascend')
+def view_as_complex(input):
+    real_part, imag_part = input.tensor_split(2, -1)
+    return _complex(real_part.squeeze(-1), imag_part.squeeze(-1))
+
+__all__.append('view_as_complex')
+
+imag_op = ops.Imag().set_device('Ascend')
+def imag(input):
+    return imag_op(input)
+
+__all__.append('imag')
+
+def glu(x, axis):
+    return pyboost_inner_prim.glu_impl(x, axis)
+
+__all__.append('glu')
+
+def ctc_loss(log_probs, targets, input_lengths, target_lengths, blank=0, reduction='mean', zero_infinity=False):
+    ctc_loss_op = _get_cache_prim(ops.CTCLossV2)(blank=blank, reduction="none", zero_infinity=zero_infinity).set_device('Ascend')
+    if targets.ndim == 1:
+        targets = targets.unsqueeze(-1)
+    loss, _ = ctc_loss_op(log_probs, targets, input_lengths, target_lengths)
+    if zero_infinity:
+        loss = select(isinf(loss), 0., loss)
+    if reduction == 'sum':
+        loss = sum_ext(loss)
+    if reduction == 'mean':
+        input_type = loss.dtype
+        target_length_t = target_lengths.clip(1., None)
+        loss = loss.astype("float32")
+        loss = div(loss, target_length_t)
+        loss = mean_ext(loss)
+        loss = loss.astype(input_type)
+    return loss
+
+__all__.append('ctc_loss')

mindnlp/core/_prims/meta.py

@@ -367,3 +367,54 @@ def sqrt(input):
     return input
 
 __all__.append('sqrt')
+
+def normal_float_float(mean, std, size, seed, offset):
+    out = Tensor_(shape=size, dtype=core.float32)
+    return core.Tensor(out)
+
+
+__all__.append('normal_float_float')
+
+def stack(tensors, dim):
+    x_shape = list(tensors[0].shape)
+    x_shape.insert(dim, len(tensors))
+    out = Tensor_(shape=tuple(x_shape), dtype=tensors[0].dtype)
+    return core.Tensor(out)
+
+__all__.append('stack')
+
+def argmax_with_value(input, dim, keepdim):
+    out_shape = list(input.shape)
+    if keepdim:
+        out_shape[dim] = 1
+    else:
+        out_shape.pop(dim)
+
+    indices = Tensor_(shape=out_shape, dtype=core.int64)
+    values = Tensor_(shape=out_shape, dtype=input.dtype)
+
+    return core.Tensor(indices), core.Tensor(values)
+
+__all__.append('argmax_with_value')
+
+def tile(input, dims):
+    input_shape = input.shape
+    out_shape = [input_shape[i] * dims[i] for i in range(input.ndim)]
+    out = Tensor_(shape=tuple(out_shape), dtype=input.dtype)
+    return core.Tensor(out)
+
+__all__.append('tile')
+
+def flatten_ext(input, start_dim, end_dim):
+    input_shape = list(input.shape)
+    if start_dim < 0:
+        start_dim = start_dim + input.ndim
+    if end_dim < 0:
+        end_dim = end_dim + input.ndim
+
+    flatten_shape = input_shape[:start_dim] + input_shape[start_dim:end_dim+1] + input_shape[end_dim+1:]
+    out = Tensor_(shape=tuple(flatten_shape), dtype=input.dtype)
+    return core.Tensor(out)
+
+__all__.append('flatten_ext')
+

mindnlp/core/_prims/numpy.py

@@ -766,3 +766,16 @@ def repeat_interleave_tensor(input, repeats, dim, _):
 
 __all__.append('repeat_interleave_tensor')
 
+def greater(input, other):
+    if not isinstance(input, numbers.Number):
+        input = input.numpy()
+    if not isinstance(other, numbers.Number):
+        other = other.numpy()
+    
+    out = input > other
+    if not isinstance(out, np.ndarray):
+        out = np.array(out)
+
+    return core.Tensor.from_numpy(out)
+
+__all__.append('greater')

mindnlp/core/_tensor.py

@@ -310,6 +310,8 @@ def __rmul__(self, other):
             return self.item() * other
         return self.__mul__(other)
 
+    def __abs__(self):
+        return ops.abs(self)
 
     def __imul__(self, other):
         return self.copy_(ops.mul(self, other))
@@ -2038,8 +2040,8 @@ def size(self, dim=None):
 
 
     # Tensor.softmax
-    def softmax(self, dim):
-        return ops.softmax(self, dim)
+    def softmax(self, dim, dtype=None):
+        return ops.softmax(self, dim, dtype=dtype)
 
     # Tensor.sort
     def sort(self, dim=-1, descending=False):
@@ -2125,7 +2127,8 @@ def sub_(self, other, *, alpha=1):
     subtract_ = sub_
 
     # Tensor.sum
-    def sum(self, dim=None, keepdim=False, dtype=None):
+    def sum(self, dim=None, keepdim=False, dtype=None, **kwargs):
+        dim = kwargs.pop('axis', dim)
         return ops.sum(self, dim, keepdim, dtype=dtype)
 
     # Tensor.sum_to_size
@@ -2155,7 +2158,8 @@ def t_(self):
         return self
 
     # Tensor.tensor_split
-
+    def tensor_split(self, indices_or_sections, dim=0):
+        return ops.tensor_split(self, indices_or_sections, dim)
 
     # Tensor.tile
     def tile(self, *dims):
@@ -2438,7 +2442,7 @@ def detach(self):
 
     # Tensor.detach_
     def detach_(self):
-        self.requires_grad_(self)
+        self.requires_grad_(False)
         return self
 
     def stub_sync(self):

mindnlp/core/nn/functional.py

@@ -197,8 +197,8 @@ def rms_norm(input, normalized_shape, weight, eps=None):
     if eps is None:
         eps = core.finfo(input.dtype).eps
     if weight is None:
-        weight = core.ones(normalized_shape)
-    return ops.rms_norm(input, weight, eps)[0]
+        weight = core.ones(normalized_shape, dtype=input.dtype, device=input.device)
+    return execute('rms_norm', input, weight, eps)
 
 def fast_gelu(x):
     return ops.fast_gelu(x)
@@ -760,7 +760,6 @@ def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
 def conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
     if isinstance(padding, str):
         return execute('conv3d_padding', input, weight, bias, stride, padding, dilation, groups)
-    print(input.device, weight.device)
     return execute('conv3d_ext', input, weight, bias, stride, padding, dilation, groups)
 
     pad_mode = 'pad'
@@ -1577,28 +1576,13 @@ def fold(input, output_size, kernel_size, dilation=1, padding=0, stride=1):
     return ops.fold(input, output_size, kernel_size, dilation, padding, stride)
 
 def ctc_loss(log_probs, targets, input_lengths, target_lengths, blank=0, reduction='mean', zero_infinity=False):
-    ctc_loss_op = _get_cache_prim(ops.CTCLossV2)(blank=blank, reduction="none", zero_infinity=zero_infinity)
-    if targets.ndim == 1:
-        targets = targets.unsqueeze(-1)
-    loss, _ = ctc_loss_op(log_probs, targets, input_lengths, target_lengths)
-    if zero_infinity:
-        loss = ops.where(ops.isinf(loss), 0., loss)
-    if reduction == 'sum':
-        loss = loss.sum()
-    if reduction == 'mean':
-        input_type = loss.dtype
-        target_length_t = target_lengths.clip(1., None)
-        loss = loss.astype("float32")
-        loss = loss / target_length_t
-        loss = loss.mean()
-        loss = loss.astype(input_type)
-    return loss
+    return execute('ctc_loss', log_probs, targets, input_lengths, target_lengths, blank, reduction, zero_infinity)
 
 def one_hot(tensor, num_classes=-1):
     return execute('one_hot_ext', tensor, num_classes)
 
 def pixel_shuffle(input, upscale_factor):
-    return ops.pixel_shuffle(input, upscale_factor)
+    return execute('pixel_shuffle', input, upscale_factor)
 
 def pixel_unshuffle(input, downscale_factor):
     return ops.pixel_unshuffle(input, downscale_factor)

mindnlp/core/nn/parameter.py

@@ -27,9 +27,6 @@ def __deepcopy__(self, memodict):
         new_obj._device = self.device
         return new_obj
 
-    def clone(self):
-        return copy.deepcopy(self)
-
     def __parameter__(self): # only for O2
         """For parse check."""
 

mindnlp/core/ops/__init__.py

@@ -1,6 +1,6 @@
 """core ops like torch funcional api"""
 from . import array, blas, comparison, pointwise, creation, random, reduction, other, \
-    tensor, _inner, optim, inplace
+    tensor, _inner, optim, inplace, complex
 from .array import *
 from .blas import *
 from .comparison import *
@@ -14,6 +14,7 @@
 from ._inner import *
 from .optim import *
 from .inplace import *
+from .complex import *
 
 def load_library(lib_path):
     raise ImportError('not support import any ops for now.')

mindnlp/core/ops/array.py

@@ -166,7 +166,12 @@ def permute(input, dims):
 def reshape(input, *shape):
     if isinstance(shape[0], (tuple, list)):
         shape = shape[0]
-    return execute("reshape", input, shape)
+    new_shape = ()
+    for s in shape:
+        if not isinstance(s, numbers.Number):
+            s = s.item()
+        new_shape += (s,)
+    return execute("reshape", input, new_shape)
 
 
 def view(input, *shape):
@@ -221,6 +226,22 @@ def split(tensor, split_size_or_sections, dim=0):
         )
     return res
 
+def split_with_sizes(input, split_sizes, dim=0):
+    assert input.dim() != 0, "split expects at least a 1-dimensional tensor"
+    dim_size = input.size(dim)
+    num_splits = len(split_sizes)
+    start_idx = 0
+
+    splits = []
+    for i in range(num_splits):
+        length = split_sizes[i]
+        assert length >= 0, f"split_with_sizes expects split_sizes have only non-negative entries, but got split_sizes={split_sizes}"
+        splits.append(
+            narrow(input, dim, start_idx, length)
+        )
+        start_idx += length
+
+    return splits
 
 # squeeze
 def squeeze(input, *dim, **kwargs):
@@ -306,10 +327,27 @@ def take_along_dim(input, indices, dim=None, *, out=None):
     return input.view(-1).gather(0, indices.view(-1))
 
 # tensor_split
-
+def tensor_split(input, indices_or_sections, dim=0):
+    if isinstance(indices_or_sections, int):
+        # 分割成大致相等的部分
+        dim_size = input.size(dim)
+        if dim_size == 0:
+            return [input] * indices_or_sections
+        split_size = (dim_size + indices_or_sections - 1) // indices_or_sections
+        return split(input, split_size, dim=dim)
+    elif isinstance(indices_or_sections, (list, tuple, core.Tensor)):
+        # 按照给定的索引分割
+        dim_size = input.size(dim)
+        indices = [0] + list(indices_or_sections) + [dim_size]
+        split_sizes = [indices[i+1] - indices[i] for i in range(len(indices)-1)]
+        return split(input, split_sizes, dim=dim)
+    else:
+        raise ValueError("indices_or_sections must be int or list/tuple of indices")
 
 # tile
 def tile(input, dims):
+    if isinstance(dims[0], (tuple, list)):
+        dims = dims[0]
     return execute("tile", input, dims)
 
 
@@ -882,6 +920,8 @@ def getitem_np(input, slice):
     return execute('getitem', input, slice)
 
 def setitem_np(input, slice, value):
+    if input.device != value.device:
+        value = value.to(input.device)
     return execute('setitem', input, slice, value)
 
 __all__ = [
@@ -926,7 +966,7 @@ def setitem_np(input, slice, value):
     "swapdims",
     "take",
     "take_along_dim",
-    # tensor_split
+    "tensor_split",
     "tile",
     "transpose",
     "unbind",
@@ -940,5 +980,6 @@ def setitem_np(input, slice, value):
     'getitem',
     'setitem',
     'getitem_np',
-    'setitem_np'
+    'setitem_np',
+    'split_with_sizes'
 ]

mindnlp/core/ops/creation.py

@@ -30,7 +30,8 @@ def frombuffer(buffer, *, dtype, count=-1, offset=0, requires_grad=False):
 
 
 # zeros
-def zeros(*size, out=None, dtype=None, layout=None, device=None, requires_grad=False):
+def zeros(*size, out=None, dtype=None, layout=None, device=None, requires_grad=False, **kwargs):
+    size = kwargs.pop('size', size)
     if dtype is None:
         dtype = get_default_dtype()
     if device is None:

mindnlp/core/ops/other.py

@@ -777,9 +777,15 @@ def unflatten(x, dim, sizes):
 
 
 # view_as_real
+def view_as_real(input):
+    real_part = input.real.unsqueeze(-1)
+    imag_part = input.imag.unsqueeze(-1)
+    return core.concat((real_part, imag_part), -1)
 
-# view_as_complex
 
+# view_as_complex
+def view_as_complex(input):
+    return execute('view_as_complex', input)
 
 # resolve_conj
 
@@ -794,6 +800,8 @@ def masked_fill(input, mask, value):
         if value == float('inf'):
             value = finfo(input.dtype).max
 
+    if isinstance(value, core.Tensor) and input.device != value.device:
+        value = value.to(input.device)
     return execute('masked_fill', input, mask, value)
 
 class finfo:
@@ -949,5 +957,7 @@ def dyn_shape(input):
     "contiguous",
     "ravel",
     "dyn_shape",
-    "diff"
+    "diff",
+    'view_as_complex',
+    'view_as_real'
 ]