fix apis for l-n class (#2157)

Author: lvyufeng

mindnlp/core/_prims/ascend.py

@@ -108,6 +108,7 @@ def pad_v3(input_x, padding, mode='constant', value=None):
     pad_op = ops.PadV3(mode=mode, paddings_contiguous=True).set_device('Ascend')
     if input_x.dtype == core.bool:
         input_x = input_x.to(core.int32)
+        value = int(value)
         out = pad_op(input_x, padding, value)
         return cast(out, core.bool)
 
@@ -117,6 +118,7 @@ def pad_v3(input_x, padding, mode='constant', value=None):
 
 __all__.append('pad_v3')
 
+
 def inplace_uniform(input, from_, to_, generator_):
     seed, offset = generator_._step(12)
     return gen_ops_prim.inplace_uniform_op(input, from_, to_, seed, offset)
@@ -413,3 +415,9 @@ def bucketize(input, boundaries, right):
     return bucketize_op(input)
 
 __all__.append('bucketize')
+
+def dropout2d(input, p):
+    dropout_2d_op = ops.Dropout2D(1.0 - p)
+    return dropout_2d_op(input)
+
+__all__.append('dropout2d')

mindnlp/core/_prims/meta.py

@@ -356,3 +356,17 @@ def bitwise_xor_tensor(input, other):
     return input
 
 __all__.append('bitwise_xor_tensor')
+
+def divmod(input, other, rounding_mode):
+    if isinstance(input, core.Tensor):
+        return input
+    return other
+
+__all__.append('divmod')
+
+def greater_equal(input, other):
+    if isinstance(input, core.Tensor):
+        return input
+    return other
+
+__all__.append('greater_equal')

mindnlp/core/_prims/numpy.py

@@ -43,17 +43,23 @@ def div(input, other):
 __all__.append('div')
 
 def pow_scalar_tensor(input, other):
-    out = np.power(input, other.numpy())
+    other = other.numpy()
+    out = np.power(input, other)
+    if out.dtype == np.float64:
+        out = out.astype(np.float32)
     return core.Tensor.from_numpy(out)
 
 __all__.append('pow_scalar_tensor')
 
 def mul(input, other):
     if not isinstance(input, numbers.Number):
         input = input.asnumpy()
-    elif not isinstance(other, numbers.Number):
+    if not isinstance(other, numbers.Number):
         other = other.asnumpy()
-    out = np.multiply(input, other)
+
+    out = input * other
+    if out.dtype == np.float64:
+        out = out.astype(np.float32)
     if not isinstance(out, np.ndarray):
         out = np.array(out)
     return core.Tensor.from_numpy(out)
@@ -598,7 +604,10 @@ def inplace_add_ext(input, other, alpha):
 __all__.append('inplace_add_ext')
 
 def pow_tensor_scalar(input, other):
-    out = np.power(input.numpy(), other)
+    input = input.numpy()
+    if input.dtype == np.int64:
+        input = input.astype(np.int32)
+    out = np.power(input, other)
     if not isinstance(out, np.ndarray):
         out = np.array(out)
     return core.Tensor.from_numpy(out)
@@ -731,8 +740,10 @@ def divmod(input, other, rounding_mode):
     if rounding_mode == 'floor':
         out = np.floor_divide(input, other)
     elif rounding_mode == 'trunc':
-        out = np.trunc(np.true_divide(input, other))
+        out = np.trunc(np.true_divide(input, other)).astype(np.int64)
 
+    if not isinstance(out, np.ndarray):
+        out = np.array(out)
     return core.Tensor.from_numpy(out)
 
 __all__.append('divmod')
@@ -801,6 +812,12 @@ def repeat_interleave_tensor(input, repeats, dim, _):
 
 __all__.append('repeat_interleave_tensor')
 
+def repeat_interleave_int(input, repeats, dim, _):
+    out = np.repeat(input.numpy(), repeats, dim)
+    return core.Tensor.from_numpy(out)
+
+__all__.append('repeat_interleave_int')
+
 def greater(input, other):
     if not isinstance(input, numbers.Number):
         input = input.numpy()
@@ -823,6 +840,8 @@ def linalg_vector_norm(input, p, dim, keepdim, dtype):
 
 def exp(input):
     out = np.exp(input.numpy())
+    if input.dtype == np.int64:
+        out = out.astype(np.float32)
     return core.Tensor.from_numpy(out)
 
 __all__.append('exp')
@@ -917,3 +936,41 @@ def floor(input):
     return core.Tensor.from_numpy(out)
 
 __all__.append('floor')
+
+def chunk(input, chunks, dim):
+    out = np.array_split(input.numpy(), chunks, dim)
+    out = [core.Tensor.from_numpy(o) for o in out]
+    return out
+
+__all__.append('chunk')
+
+def narrow(input, dim, start, length):
+    slices = [slice(None)] * input.ndim
+    # 将指定维度的切片修改为 [start: start+length]
+    slices[dim] = slice(start, start + length)
+    # 应用切片并返回视图
+    out = input.numpy()[tuple(slices)]
+    return core.Tensor.from_numpy(out)
+
+__all__.append('narrow')
+
+def roll(input, shifts, dims):
+    out = np.roll(input.numpy(), shifts, dims)
+    return core.Tensor.from_numpy(out)
+
+__all__.append('roll')
+
+def outer(input, other):
+    out = np.outer(input.numpy(), other.numpy())
+    return core.Tensor.from_numpy(out)
+
+__all__.append('outer')
+
+def one_hot_ext(tensor, num_classes=-1):
+    if num_classes == -1:
+        num_classes = np.max(tensor.numpy()) + 1  # 自动确定类别数[2](@ref)
+    
+    out = np.eye(num_classes)[tensor.numpy()]
+    return core.Tensor.from_numpy(out)
+
+__all__.append('one_hot_ext')

mindnlp/core/_tensor.py

@@ -173,6 +173,7 @@ def cuda(self, device=None, non_blocking=False):
 
     def requires_grad_(self, requires_grad=True):
         self.requires_grad = requires_grad
+        return self
 
     def __reduce_ex__(self, protocol):
         if isinstance(self, StubTensor):
@@ -293,6 +294,8 @@ def __rtruediv__ (self, other):
         return ops.div(other, self)
 
     def __ne__(self, other):
+        if isinstance(other, list):
+            return True
         return ops.ne(self, other)
 
     def __neg__(self):
@@ -2122,10 +2125,10 @@ def untyped_storage(self):
 
 
     # Tensor.stride
-    def stride(self, dim=None):
-        if dim is None:
-            return self._data.stride()
-        return self._data.stride()[dim]
+    # def stride(self, dim=None):
+    #     if dim is None:
+    #         return self.stride()
+    #     return self.stride()[dim]
 
 
     # Tensor.sub

mindnlp/core/nn/functional.py

@@ -7,7 +7,7 @@
 from mindnlp import core
 from mindnlp.core.executor import execute
 
-from ..configs import DEVICE_TARGET, ON_ORANGE_PI, use_pyboost, ON_A1, ON_A2
+from ..configs import ON_ORANGE_PI, use_pyboost, ON_A1, ON_A2
 
 generator_step_ = 12
 
@@ -74,7 +74,7 @@ def hardsigmoid(input, inplace=False):
     return ops.hardsigmoid(input)
 
 def hardswish(input: core.Tensor, inplace: bool = False) -> core.Tensor:
-    return ops.hardswish(input)
+    return execute('hswish', input)
 
 def hardshrink(input, lambd=0.5):
     return execute('hard_shrink', input, lambd)
@@ -129,7 +129,7 @@ def adaptive_avg_pool2d(input, output_size):
     return execute('adaptive_avg_pool2d_ext', input, output_size)
 
 def dropout(input, p=0.5, training=True, inplace=False):
-    if not training:
+    if not training or p==0:
         return input
     out, _ = execute('dropout_ext', input, p)
     if inplace:
@@ -138,7 +138,10 @@ def dropout(input, p=0.5, training=True, inplace=False):
     return out
 
 def dropout2d(input, p=0.5, training=False):
-    return ops.dropout2d(input, p, training)
+    if not training or p==0:
+        return input
+    out, _ = execute('dropout2d', input, p)
+    return out
 
 def drop_and_mask(keep_prob, seed=None):
     seed0, seed1 = _get_seed(seed, "dropout")
@@ -301,6 +304,9 @@ def pad(input, pad, mode='constant', value=None):
             return execute('pad_v3', input, new_pad, mode)
         if value is None:
             value = 0
+        if mode == "replicate":
+            mode = "edge"
+            return execute('pad_v3', input, new_pad, mode)
         return execute('pad_v3', input, new_pad, mode, value)
     out = input
     if (isinstance(pad, tuple) and not pad):
@@ -1541,8 +1547,8 @@ def _canonical_mask(
 ) -> Optional[core.Tensor]:
     if mask is not None:
         _mask_dtype = mask.dtype
-        _mask_is_float = ops.is_floating_point(mask)
-        if _mask_dtype != mindspore.bool_ and not _mask_is_float:
+        _mask_is_float = core.is_floating_point(mask)
+        if _mask_dtype != core.bool and not _mask_is_float:
             raise AssertionError(
                 f"only bool and floating types of {mask_name} are supported")
         if check_other and other_type is not None:
@@ -1552,8 +1558,8 @@ def _canonical_mask(
                     "is deprecated. Use same type for both instead."
                 )
         if not _mask_is_float:
-            zero_tensor = ops.zeros_like(mask, dtype=target_type)
-            mask = ops.where(mask, core.Tensor(float("-inf"), target_type), zero_tensor)
+            zero_tensor = core.zeros_like(mask, dtype=target_type, device=mask.device)
+            mask = core.where(mask, core.tensor(float("-inf"), dtype=target_type, device=mask.device), zero_tensor)
             # mask = (
             #     ops.zeros_like(mask, dtype=target_type)
             #     .masked_fill_(mask, float("-inf"))
@@ -1571,14 +1577,9 @@ def unfold(input, kernel_size, dilation=1, padding=0, stride=1):
     if ON_A1:
         return execute('im2col', input, kernel_size, dilation, padding, stride)
     return execute('im2col_ext', input, kernel_size, dilation, padding, stride)
-    if use_pyboost() and not ON_A1:
-        return mint.nn.functional.unfold(input, kernel_size, dilation, padding, stride)
-    return ops.unfold(input, kernel_size, dilation, padding, stride)
 
 def fold(input, output_size, kernel_size, dilation=1, padding=0, stride=1):
-    if use_pyboost():
-        return mint.nn.functional.fold(input, output_size, kernel_size, dilation, padding, stride)
-    return ops.fold(input, output_size, kernel_size, dilation, padding, stride)
+    return execute('col2im_ext', 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):
     return execute('ctc_loss', log_probs, targets, input_lengths, target_lengths, blank, reduction, zero_infinity)

mindnlp/core/npu/__init__.py

@@ -115,7 +115,8 @@ def npu_fusion_attention(query, key, value, head_num, input_layout, *, pse=None,
                          scale=1., keep_prob=1., pre_tockens=2147483647, next_tockens=2147483647, inner_precise=0,
                          drop_mask=None, prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0,
                          gen_mask_parallel=True, sync=False, pse_type=1, q_start_idx=None, kv_start_idx=None):
-    output = gen.flash_attention_score_impl(
+    output = execute(
+        'flash_attention_score',
         query, key, value, real_shift=pse, padding_mask=padding_mask, drop_mask=drop_mask,
         attn_mask=atten_mask, prefix=prefix, actual_seq_qlen=actual_seq_qlen,
         actual_seq_kvlen=actual_seq_kvlen, head_num=head_num, keep_prob=float(keep_prob),

mindnlp/core/ops/array.py

@@ -201,6 +201,7 @@ def _get_moved_perm(ndim, source, destination):
 
 # narrow
 def narrow(input, dim, start, length):
+    length = length.item() if not isinstance(length, int) else length
     return execute("narrow", input, dim, start, length)
 
 
@@ -393,7 +394,7 @@ def tensor_split(input, indices_or_sections, dim=0):
 def tile(input, dims):
     if isinstance(dims[0], (tuple, list)):
         dims = dims[0]
-    return execute("tile", input, dims)
+    return execute("tile", input, tuple(dims))
 
 
 # transpose

mindnlp/core/ops/creation.py

@@ -12,7 +12,10 @@
 from .._bind import get_default_dtype, get_device_in_context
 
 def as_strided(self, size, stride, storage_offset=None):
-    return execute('as_strided', self, size, stride, storage_offset)
+    size = [s if isinstance(s, int) else s.item() for s in size]
+    if storage_offset is None:
+        storage_offset = 0
+    return execute('as_strided', self, tuple(size), tuple(stride), storage_offset)
 
 # from_numpy
 def from_numpy(ndarray):
@@ -37,7 +40,7 @@ def zeros(*size, out=None, dtype=None, layout=None, device=None, requires_grad=F
     if device is None:
         device = get_device_in_context()
 
-    if isinstance(device, str):
+    if isinstance(device, (str, int)):
         device = core.device(device)
     if len(size) > 0 and isinstance(size[0], (tuple, list)):
         size = size[0]

mindnlp/core/ops/other.py

@@ -97,7 +97,8 @@ def clone(input, *, memory_format=core.preserve_format):
 # cumprod
 
 # cumsum
-def cumsum(input, dim, dtype=None):
+def cumsum(input, dim=None, dtype=None, **kwargs):
+    dim = kwargs.pop('axis', dim)
     if input.dtype in [core.int64, core.bool]:
         return execute('cumsum_ext', input.int(), dim, None).long()
     if dtype is not None and dtype == core.int64:

mindnlp/core/ops/reduction.py

@@ -40,7 +40,8 @@ def any(input, dim=None, keepdim=False):
     return execute('reduce_any', input, dim, keepdim)
 
 # max
-def max(input, dim=None, keepdim=False, *, out=None):
+def max(input, dim=None, keepdim=False, *, out=None, **kwargs):
+    dim = kwargs.pop('axis', dim)
     if dim is None and not keepdim:
         return execute('max', input)
     if core.is_tensor(dim):