add cumsum to _legacy/functional and longformer_embedding

mindspore-lab · lvyufeng · May 13, 2023 · Mar 8, 2023 · Mar 8, 2023 · Mar 8, 2023
commit ae55a61f2cb1d99f42ee1ad0fa55c02e807e0a25
diff --git a/mindnlp/_legacy/functional.py b/mindnlp/_legacy/functional.py
@@ -26,11 +26,13 @@
 from mindspore.common import dtype as mstype
 from mindspore.ops._primitive_cache import _get_cache_prim
 from packaging import version
+from mindspore.ops import operations as P
 
 MS_COMPATIBLE_VERSION = '1.10.1'
 
 cast_ = ops.Cast()
 scalar_to_tensor_ = ops.ScalarToTensor()
+cumsum_ = P.CumSum()
 
 def kl_div(inputs, target, reduction='none', log_target=False):
     """KLDiv function."""
@@ -642,3 +644,57 @@ def arange(start=0, end=None, step=1, *, dtype=None):
     if dtype is not None:
         data = cast_(data, dtype)
     return data
+
+def cumsum(x, axis, dtype=None):
+    """
+    Computes the cumulative sum of input Tensor along `axis`.
+
+    .. math::
+
+        y_i = x_1 + x_2 + x_3 + ... + x_i
+
+    Note:
+        On Ascend, the dtype of `x` only support :int8, uint8, int32, float16 or float32 in case of static shape.
+        For the case of dynamic shape, the dtype of `x` only support int32, float16 or float32.
+
+    Args:
+        x (Tensor): The input Tensor to accumulate.
+        axis (int): Axis along which the cumulative sum is computed.
+        dtype (:class:`mindspore.dtype`, optional): The desired dtype of returned Tensor. If specified,
+            the input Tensor will be cast to `dtype` before the computation. This is useful for preventing overflows.
+            If not specified, stay the same as original Tensor. Default: None.
+
+    Returns:
+        Tensor, the shape of the output Tensor is consistent with the input Tensor's.
+
+    Raises:
+        TypeError: If `x` is not a Tensor.
+        ValueError: If the axis is out of range.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor
+        >>> import mindspore.ops as ops
+        >>> x = Tensor(np.array([[3, 4, 6, 10], [1, 6, 7, 9], [4, 3, 8, 7], [1, 3, 7, 9]]).astype(np.float32))
+        >>> # case 1: along the axis 0
+        >>> y = ops.cumsum(x, 0)
+        >>> print(y)
+        [[ 3.  4.  6. 10.]
+         [ 4. 10. 13. 19.]
+         [ 8. 13. 21. 26.]
+         [ 9. 16. 28. 35.]]
+        >>> # case 2: along the axis 1
+        >>> y = ops.cumsum(x, 1)
+        >>> print(y)
+        [[ 3.  7. 13. 23.]
+         [ 1.  7. 14. 23.]
+         [ 4.  7. 15. 22.]
+         [ 1.  4. 11. 20.]]
+    """
+    if dtype is not None and x.dtype != dtype:
+        x = x.astype(dtype, copy=False)
+    return cumsum_(x, axis)
diff --git a/mindnlp/models/longformer/longformer.py b/mindnlp/models/longformer/longformer.py
@@ -23,7 +23,7 @@
 # pylint:
 import mindspore
 from mindspore import nn
-
+from ..._legacy import functional
 activation_map = {
     'relu': nn.ReLU(),
     'gelu': nn.GELU(False),
@@ -83,7 +83,7 @@ def create_position_ids_from_input_ids(input_ids, padding_idx):
     """
     # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
     mask = mindspore.ops.not_equal(input_ids, padding_idx).astype(mindspore.int32)
-    incremental_indices = mindspore.ops.cumsum(mask, axis=1, dtype=mask.dtype) * mask
+    incremental_indices = functional.cumsum(mask, axis=1, dtype=mask.dtype) * mask
     return incremental_indices.long() + padding_idx