fix apis for e class (#2152)

Author: lvyufeng

mindnlp/core/_bind.py

@@ -31,7 +31,7 @@ def get_autocast_dtype(device_type):
 def get_autocast_gpu_dtype():
     return AUTO_CAST_DTYE['cuda']
 
-def is_autocast_enabled(device):
+def is_autocast_enabled(device=None):
     return AUTO_CAST_ENABLED
 
 def set_default_dtype(dtype):

mindnlp/core/_prims/ascend.py

@@ -375,3 +375,24 @@ def ctc_loss(log_probs, targets, input_lengths, target_lengths, blank=0, reducti
     return loss
 
 __all__.append('ctc_loss')
+
+def reduce_max(input, dim, keepdim):
+    return pyboost_inner_prim.reduce_max_impl(input, dim, keepdim)
+
+__all__.append('reduce_max')
+
+def elu(input, alpha):
+    return pyboost_inner_prim.elu_ext_impl(input, alpha)
+
+__all__.append('elu')
+
+dynamic_rnn_op = ops.DynamicRNN().set_device('Ascend')
+def dynamic_rnn(*args):
+    return dynamic_rnn_op(*args)
+
+__all__.append('dynamic_rnn')
+
+def cross(input, other, dim=None, *, out=None):
+    return pyboost_inner_prim.cross_impl(input, other, dim)
+
+__all__.append('cross')

mindnlp/core/_prims/meta.py

@@ -418,3 +418,22 @@ def flatten_ext(input, start_dim, end_dim):
 
 __all__.append('flatten_ext')
 
+def cumsum_ext(input, dim, dtype):
+    return input
+
+__all__.append('cumsum_ext')
+
+def squeeze(input, dim):
+    input_shape = list(input.shape)
+    if isinstance(dim, int):
+        dim = (dim,)
+    
+    new_shape = ()
+    for idx, s in enumerate(input_shape):
+        if idx not in dim and s != 1:
+            new_shape += (s,)
+
+    out = Tensor_(shape=tuple(new_shape), dtype=input.dtype)
+    return core.Tensor(out)
+
+__all__.append('squeeze')

mindnlp/core/_prims/numpy.py

@@ -779,3 +779,9 @@ def greater(input, other):
     return core.Tensor.from_numpy(out)
 
 __all__.append('greater')
+
+def linalg_vector_norm(input, p, dim, keepdim, dtype):
+    out = np.linalg.norm(input.numpy(), p, dim, keepdim)
+    return core.Tensor.from_numpy(out)
+
+__all__.append('linalg_vector_norm')

mindnlp/core/_tensor.py

@@ -386,11 +386,20 @@ def new(self, *shape):
 
     # Tensor.new_tensor
     def new_tensor(self, data, *, dtype=None, device=None, requires_grad=False, layout=None, pin_memory=False):
-        return tensor(data, dtype=dtype if dtype is not None else self.dtype)
+        if device is None:
+            device = self.device
+        if dtype is None:
+            dtype = self.dtype
+
+        return tensor(data, dtype=dtype, device=device)
 
     # Tensor.new_full
     def new_full(self, size, fill_value, *, dtype=None, device=None, requires_grad=False, layout=None, pin_memory=False):
-        return ops.full(size, fill_value, dtype=dtype if dtype is not None else self.dtype)
+        if device is None:
+            device = self.device
+        if dtype is None:
+            dtype = self.dtype
+        return ops.full(size, fill_value, dtype=dtype, device=device)
 
     # Tensor.new_empty
     def new_empty(self, size, *, dtype=None, device=None, requires_grad=False, layout=None, pin_memory=False):
@@ -1058,7 +1067,8 @@ def expm1_(self):
 
 
     # Tensor.expand
-    def expand(self, *size):
+    def expand(self, *size, **kwargs):
+        size = kwargs.pop('size', size)
         if len(size) == 1:
             size = size[0]
         return self.broadcast_to(size)
@@ -1284,7 +1294,7 @@ def index_select(self, dim, index):
 
     # Tensor.int
     def int(self):
-        return self.to(mindspore.int64)
+        return self.to(mindspore.int32)
 
     # Tensor.int_repr
 
@@ -2129,6 +2139,7 @@ def sub_(self, other, *, alpha=1):
     # Tensor.sum
     def sum(self, dim=None, keepdim=False, dtype=None, **kwargs):
         dim = kwargs.pop('axis', dim)
+        keepdim = kwargs.pop('keepdims', keepdim)
         return ops.sum(self, dim, keepdim, dtype=dtype)
 
     # Tensor.sum_to_size
@@ -2551,6 +2562,9 @@ def log_softmax(self, dim):
     def char(self):
         return self.to(core.int8)
 
+    def cross(self, other, dim=None):
+        return ops.cross(self, other, dim)
+
     @property
     def is_nested(self):
         return False

mindnlp/core/linalg/__init__.py

@@ -25,7 +25,7 @@ def cholesky_ex(A, *, upper=False, check_errors=False, out=None):
 
 
 def norm(A, ord=None, dim=None, keepdim=False, *, out=None, dtype=None):
-    return mint.norm(A, 2 if ord is None else ord, dim, keepdim, dtype=dtype)
+    return core.norm(A, 2 if ord is None else ord, dim, keepdim, dtype=dtype)
 
 def vector_norm(x, ord=2, dim=None, keepdim=False, *, dtype=None, out=None):
     return execute('linalg_vector_norm', x, ord, dim, keepdim, dtype=dtype)

mindnlp/core/nn/functional.py

@@ -47,7 +47,7 @@ def relu6(input):
     return execute('relu6', input)
 
 def elu(input, alpha=1.0):
-    return execute('relu6', input, alpha)
+    return execute('elu', input, alpha)
 
 def glu(input, dim=-1):
     return execute('glu', input, dim)
@@ -59,9 +59,7 @@ def logsigmoid(input):
     return execute('logsigmoid', input)
 
 def leaky_relu(input, alpha=0.2):
-    if use_pyboost():
-        return mint.nn.functional.leaky_relu(input, alpha)
-    return ops.leaky_relu(input, alpha)
+    return execute('leaky_relu_ext', input, alpha)
 
 def prelu(input, weight):
     return execute('prelu', input, weight)
@@ -284,6 +282,9 @@ def _circular_pad(input_x, padding):
     return out
 
 def pad(input, pad, mode='constant', value=None):
+    if isinstance(pad, tuple):
+        pad = tuple(p if isinstance(p, int) else p.item() for p in pad)
+
     if input.device.type in ['cpu', 'meta'] or ON_A1:
         new_pad = ()
         for idx, pad_v in enumerate(pad):
@@ -296,6 +297,8 @@ def pad(input, pad, mode='constant', value=None):
             return input
         if mode == 'circular':
             return custom_circular_pad(input, pad)
+        elif mode == 'reflect':
+            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):

mindnlp/core/nn/modules/rnn.py

@@ -19,6 +19,7 @@
 from mindspore import ops as P
 
 from mindnlp import core
+from mindnlp.core.executor import execute
 from .module import Module
 from .dropout import Dropout
 from ..parameter import Parameter
@@ -29,9 +30,9 @@
 __all__ = ['LSTM', 'GRU', 'RNN']
 
 
-def _init_state(shape, dtype, is_lstm):
-    hx = ops.zeros(*shape, dtype=dtype)
-    cx = ops.zeros(*shape, dtype=dtype)
+def _init_state(shape, dtype, device, is_lstm):
+    hx = ops.zeros(*shape, dtype=dtype, device=device)
+    cx = ops.zeros(*shape, dtype=dtype, device=device)
     if is_lstm:
         return (hx, cx)
     return hx
@@ -285,7 +286,7 @@ def forward(self, x, h_0, seq_length, w_ih, w_hh, b_ih, b_hh):
         w_hh = ops.cat((w_hh_i, w_hh_g, w_hh_f, w_hh_o), 0)
         weight = ops.cat((w_ih, w_hh), 1)
         if b_ih is None:
-            bias = ops.zeros(w_ih.shape[0], dtype=w_ih.dtype)
+            bias = ops.zeros(w_ih.shape[0], dtype=w_ih.dtype, device=w_ih.device)
         else:
             b_ih_i, b_ih_f, b_ih_g, b_ih_o = ops.chunk(b_ih, 4, 0)
             b_hh_i, b_hh_f, b_hh_g, b_hh_o = ops.chunk(b_hh, 4, 0)
@@ -294,7 +295,8 @@ def forward(self, x, h_0, seq_length, w_ih, w_hh, b_ih, b_hh):
                             b_ih_f + b_hh_f, \
                             b_ih_o + b_hh_o), 0)
 
-        outputs, h, c, _, _, _, _, _ = self.lstm(x.to(core.float16), \
+        outputs, h, c, _, _, _, _, _ = execute('dynamic_rnn',
+                                                 x.to(core.float16), \
                                                  ops.transpose(weight, 1, 0).to(core.float16), \
                                                  bias.to(core.float16), None, \
                                                  h_0[0].unsqueeze(0).to(core.float16), \
@@ -314,8 +316,8 @@ class _RNNBase(Module):
     '''Basic class for RNN operators'''
 
     def __init__(self, mode, input_size, hidden_size, num_layers=1, bias=True,
-                 batch_first=False, dropout=0., bidirectional=False, dtype=None):
-        factory_kwargs = {'dtype': dtype}
+                 batch_first=False, dropout=0., bidirectional=False, dtype=None, device=None):
+        factory_kwargs = {'dtype': dtype, 'device': device}
         super().__init__()
 
         if not 0 <= dropout < 1:
@@ -495,7 +497,7 @@ def forward(self, x, hx=None, seq_length=None):
         x_dtype = x.dtype
         if hx is None:
             hx = _init_state((self.num_layers * num_directions, max_batch_size, self.hidden_size), \
-                             x_dtype, self.is_lstm)
+                             x_dtype, x.device, self.is_lstm)
         if self.batch_first:
             x = ops.permute(x, (1, 0, 2))
         if self.bidirectional:

mindnlp/core/ops/array.py

@@ -218,6 +218,7 @@ def split(tensor, split_size_or_sections, dim=0):
     if isinstance(split_size_or_sections, int):
         res = execute("split_tensor", tensor, split_size_or_sections, dim)
     elif isinstance(split_size_or_sections, (list, tuple)):
+        split_size_or_sections = tuple(s.item() if isinstance(s, core.Tensor) else s for s in split_size_or_sections)
         res = execute("split_with_size", tensor, split_size_or_sections, dim)
     else:
         raise TypeError(
@@ -227,21 +228,7 @@ 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
+    return execute("split_with_size", input, split_sizes, dim)
 
 # squeeze
 def squeeze(input, *dim, **kwargs):

mindnlp/core/ops/other.py

@@ -88,6 +88,10 @@ def clone(input, *, memory_format=core.preserve_format):
 
 # cumsum
 def cumsum(input, dim, dtype=None):
+    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:
+        return execute('cumsum_ext', input, dim, None).long()
     return execute('cumsum_ext', input, dim, dtype)
 
 # diag
@@ -928,6 +932,12 @@ def contiguous(input):
 def dyn_shape(input):
     return execute('dyn_shape', input)
 
+def cross(input, other, dim=None, *, out=None):
+    if dim is None:
+        dim = -65530
+    return execute('cross', input, other, dim)
+
+
 __all__ = [
     "bincount",
     "broadcast_shapes",
@@ -936,6 +946,7 @@ def dyn_shape(input):
     "cdist",
     "clone",
     "contains",
+    "cross",
     "cumsum",
     "diag",
     "diagonal",