mindspore 2.4.1__cp310-cp310-manylinux1_x86_64.whl → 2.5.0__cp310-cp310-manylinux1_x86_64.whl

mindspore/ops/composite/multitype_ops/_compile_utils.py

@@ -16,8 +16,9 @@
 """constexpr util"""
 from __future__ import absolute_import
 from enum import IntEnum
+import numpy as np
 
-
+from mindspore._c_expression import Tensor as Tensor_
 from mindspore.ops.composite.multitype_ops import _constexpr_utils as const_utils
 from mindspore.ops import functional as F
 from mindspore.ops import operations as P
@@ -35,6 +36,8 @@ from mindspore import ops
 from mindspore.ops.primitive import _primexpr
 from mindspore import _checkparam as validator
 from mindspore.common._stub_tensor import _convert_stub
+from mindspore.ops.auto_generate.gen_ops_prim import select_ext_op, slice_ext_op, inplace_copy_op, \
+    index_op, inplace_index_put_op
 
 slice_get_item = SliceGetItem()
 hyper_map = base.HyperMap()
@@ -45,9 +48,15 @@ is_parameter = IsParameter()
 getitem_tensor_index_info = GetitemTensorIndexInfo(const_utils.is_ascend())
 setitem_tensor_index_info = SetitemTensorIndexInfo(const_utils.is_ascend())
 
-selevt_view = SelectView()
+select_view = SelectView()
 copy_with_slice = CopyWithSlice()
 
+tensor_1d = Tensor([0], dtype=mstype.int64)
+empty_tensor_1d = Tensor_(shape=(0,), dtype=mstype.int64)
+empty_tensor_9d = Tensor_(shape=(0,)*9, dtype=mstype.int64)
+
+
+
 def strided_slice(data, begin_strides, end_strides, step_strides, begin_mask=0, end_mask=0, ellipsis_mask=0,
                   new_axis_mask=0, shrink_axis_mask=0):
     """strided_slice primitive cache"""
@@ -148,7 +157,7 @@ def data_update_by_ops(transfer_type, arg, data, new_index, origin_data, value=N
     elif transfer_type == ValueTransferType.kSelect:
         data = F.select(Tensor(new_index), value, data)
     elif transfer_type == ValueTransferType.kSelectView:
-        data = selevt_view(data, arg[0], arg[1])
+        data = select_view(data, arg[0], arg[1])
     elif transfer_type == ValueTransferType.kCopyView:
         value = _broadcast(F.shape(data), F.cast(value, F.dtype(data)))
         data = copy_with_slice(data, value)
@@ -196,14 +205,14 @@ def value_update(transfer_types, args, data, value):
     return value
 
 
-def _tensor_getitem(self, index):
+def _tensor_getitem_origin(self, index):
     """Handle tensor getitem"""
     new_index, tensor_update_types, tensor_update_args = getitem_tensor_index_info(
         self, index)
     return data_update(tensor_update_types, tensor_update_args, self, new_index)
 
 
-def _tensor_setitem(self, index, value):
+def _tensor_setitem_origin(self, index, value):
     """Handle tensor setitem"""
     setitem_info = setitem_tensor_index_info(self, index, value)
     new_index = setitem_info[0]
@@ -218,8 +227,213 @@ def _tensor_setitem(self, index, value):
     return output
 
 
-setattr(tensor_operator_registry, "__getitem__", _tensor_getitem)
-setattr(tensor_operator_registry, "__setitem__", _tensor_setitem)
+setattr(tensor_operator_registry, "_tensor_getitem_origin", _tensor_getitem_origin)
+setattr(tensor_operator_registry, "_tensor_setitem_origin", _tensor_setitem_origin)
+
+
+def _record_tensor_index(index, remain_indexes, dim):
+    """Record indexes remained to be used by aclnnIndex/aclnnIndexPut"""
+    if len(remain_indexes) > dim:
+        remain_indexes[dim] = index
+        return remain_indexes
+
+    while dim > len(remain_indexes):
+        # use empty_tensor with dim_num 9 to indicate unused dim
+        remain_indexes.append(empty_tensor_9d)
+
+    remain_indexes.append(index)
+    return remain_indexes
+
+
+def _count_indexed_dims(indexes):
+    """Count indexed dims"""
+    count = 0
+    for index in indexes:
+        if isinstance(index, Tensor):
+            if index.dtype == mstype.bool_:
+                count += index.ndim
+            else:
+                count += 1
+        elif not isinstance(index, (type(None), type(...), bool)):
+            count += 1
+    return count
+
+
+def _do_select(self: Tensor, dim: int, index: int, dim_index: int, self_shape: list):
+    """call select view operator"""
+    if not self_shape:
+        raise TypeError("Invalid index of a 0-dim tensor.")
+    dim_size = self_shape[dim]
+    if index >= dim_size or index < -dim_size:
+        raise IndexError(f"Index {index} is out of bounds for dimension {dim_index} with size {dim_size}")
+    index = index + dim_size if index < 0 else index
+    return select_ext_op(self, dim, index)
+
+
+def _do_slice(self: Tensor, dim: int, index: slice, self_shape: list):
+    """call slice view operator"""
+    def _get_index(index, default):
+        if index is None:
+            return default
+        if isinstance(index, Tensor):
+            return index.__index__()
+        return index
+
+    if not self_shape:
+        raise TypeError("Invalid index of a 0-dim tensor.")
+    step = _get_index(index.step, 1)
+    if step <= 0:
+        raise ValueError("slice step must be positive")
+    start = _get_index(index.start, 0)
+    end = _get_index(index.stop, self_shape[dim])
+    if start == 0 and end == self_shape[dim] and step == 1:
+        return self
+    return slice_ext_op(self, dim, start, end, step)
+
+
+def _process_dim_in_multi_dim_index(prev_result, orig_tensor, index, dim, indexed_dims, dim_index, remain_indexes,
+                                    prev_shape):
+    """Process dim in multi dim index"""
+    if isinstance(index, bool):
+        result = F.expand_dims(prev_result, dim)
+        index_for_bool = tensor_1d if index else empty_tensor_1d
+        _record_tensor_index(index_for_bool, remain_indexes, dim)
+        prev_shape.insert(dim, 1)
+        dim += 1
+        return result, dim, remain_indexes, prev_shape
+    if isinstance(index, int):
+        result = _do_select(prev_result, dim, index, dim_index, prev_shape)
+        del prev_shape[dim]
+        return result, dim, remain_indexes, prev_shape
+    if isinstance(index, slice):
+        result = _do_slice(prev_result, dim, index, prev_shape)
+        # current dim in prev_shape will not be used later, ignore it
+        dim += 1
+        return result, dim, remain_indexes, prev_shape
+    if isinstance(index, type(...)):
+        dim += (orig_tensor.ndim - indexed_dims)
+        return prev_result, dim, remain_indexes, prev_shape
+    if index is None:
+        result = F.expand_dims(prev_result, dim)
+        prev_shape.insert(dim, 1)
+        dim += 1
+        return result, dim, remain_indexes, prev_shape
+    if isinstance(index, Tensor):
+        result = prev_result
+        if index.ndim == 0 and index.dtype in mstype.int_type + mstype.uint_type + (mstype.bool_,):
+            if index.dtype in mstype.int_type + mstype.uint_type:
+                result = _do_select(prev_result, dim, index.item(), dim_index, prev_shape)
+                del prev_shape[dim]
+                return result, dim, remain_indexes, prev_shape
+            # process index with Tensor bool type
+            result = F.expand_dims(prev_result, dim)
+            index_for_bool = tensor_1d if index else empty_tensor_1d
+            _record_tensor_index(index_for_bool, remain_indexes, dim)
+            prev_shape.insert(dim, 1)
+            dim += 1
+            return result, dim, remain_indexes, prev_shape
+        _record_tensor_index(index, remain_indexes, dim)
+        dim += 1
+        return result, dim, remain_indexes, prev_shape
+    raise IndexError(f"Invalid tensor index type {index}")
+
+
+def _process_multi_dim_index(self, indexes, remain_indexes, indexed_dims):
+    """Process indexes in tuple"""
+    self_viewed = self
+    self_viewed_shape = list(self.shape)
+    dim = 0
+    for i, index in enumerate(indexes):
+        if isinstance(index, (list, tuple, np.ndarray)):
+            index_np = np.array(index) if isinstance(index, (list, tuple)) else index
+            if index_np.dtype in (np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64,
+                                  np.float16, np.float32, np.float64):
+                index = Tensor(index_np, mstype.int64)
+            elif index_np.dtype == np.bool_:
+                index = Tensor(index_np, mstype.bool_)
+            else:
+                raise TypeError(f"Index {index} contain unsupported elements")
+        self_viewed, dim, remain_indexes, self_viewed_shape = _process_dim_in_multi_dim_index(
+            self_viewed, self, index, dim, indexed_dims, i, remain_indexes, self_viewed_shape)
+    return self_viewed, remain_indexes
+
+
+def _wrap_index_to_tuple(index):
+    """Wrap index to tuple"""
+    if isinstance(index, tuple):
+        return index
+    if isinstance(index, list):
+        if len(index) < 32 and any(isinstance(i, (Tensor, list, tuple, slice, type(None), type(...))) for i in index):
+            return tuple(index)
+    return (index,)
+
+
+def _tensor_getitem(self, index):
+    """Handle tensor getitem"""
+    if isinstance(index, bool):
+        self_viewed = F.expand_dims(self, 0)
+        index_for_bool = tensor_1d if index else empty_tensor_1d
+        return index_op(self_viewed, [index_for_bool])
+    if isinstance(index, int):
+        return _do_select(self, 0, index, 0, list(self.shape))
+    if isinstance(index, slice):
+        result = _do_slice(self, 0, index, list(self.shape))
+        return result
+    if index is None:
+        return F.expand_dims(self, 0)
+    if isinstance(index, type(...)):
+        return self
+    indexes = _wrap_index_to_tuple(index)
+    indexed_dims = _count_indexed_dims(indexes)
+    if self.ndim < indexed_dims:
+        raise IndexError(f"too many indices for tensor with dimension size {self.ndim}")
+    remain_indexes = []
+    self_viewed, remain_indexes = _process_multi_dim_index(self, indexes, remain_indexes, indexed_dims)
+    if not remain_indexes:
+        return self_viewed
+    return index_op(self_viewed, remain_indexes)
+
+
+def _tensor_setitem(self, index, value):
+    """Handle tensor setitem"""
+    if not isinstance(value, Tensor):
+        if isinstance(value, (bool, int, float)):
+            value = Tensor(value, dtype=self.dtype)
+        else:
+            raise TypeError(f"Can't assign a {type(value)} to a {self.dtype}.")
+
+    if isinstance(index, bool) and index is False:
+        return self
+    if isinstance(index, type(...)):
+        inplace_copy_op(self, value)
+        return self
+    if index is None or (isinstance(index, bool) and index is True):
+        self_viewed = F.expand_dims(self, 0)
+        inplace_copy_op(self_viewed, value)
+        return self
+    if isinstance(index, int):
+        self_viewed = _do_select(self, 0, index, 0, list(self.shape))
+        inplace_copy_op(self_viewed, value)
+        return self
+    if isinstance(index, slice):
+        self_viewed = _do_slice(self, 0, index, list(self.shape))
+        inplace_copy_op(self_viewed, value)
+        return self
+    indexes = _wrap_index_to_tuple(index)
+    indexed_dims = _count_indexed_dims(indexes)
+    if self.ndim < indexed_dims:
+        raise IndexError(f"too many indices for tensor with dimension size {self.ndim}")
+    remain_indexes = []
+    self_viewed, remain_indexes = _process_multi_dim_index(self, indexes, remain_indexes, indexed_dims)
+    if not remain_indexes:
+        inplace_copy_op(self_viewed, value)
+        return self
+    inplace_index_put_op(self_viewed, remain_indexes, value)
+    return self
+
+
+setattr(tensor_operator_registry, "_tensor_getitem", _tensor_getitem)
+setattr(tensor_operator_registry, "_tensor_setitem", _tensor_setitem)
 
 
 def _tensor_add(self, other):
@@ -313,31 +527,16 @@ def _check_scalar_tensor_args(args):
         const_utils.raise_value_error("For item, the index of scalar Tensor should not be set.")
 
 
-def tensor_item(data, *args):
-    """Tensor getitem by index whose dtype is int or tuple with int."""
-    # transform a.item(tuple(int)) -> a.item(int1,int2...intN)
+def tensor_item(data):
+    """Tensor getitem which has only one element."""
     if data.ndim == 0:
-        _check_scalar_tensor_args(args)
         return TensorToScalar()(data)
-    if len(args) == 1 and isinstance(args[0], tuple):
-        args = args[0]
-
-    args_types = hyper_map(F.typeof, args)
-    if not args or const_utils.judge_index_type(args_types[0], mstype.type_none):
-        if data.shape == (1,):
-            return TensorToScalar()(data[0])
-        const_utils.raise_value_error("Can only convert an array of size 1 to a Python scalar")
-
-    if not const_utils.judge_indexes_types(args_types, mstype.int64):
-        const_utils.raise_type_error("The index object cannot be interpreted as an integer")
-
-    if len(args) == data.ndim:
-        return tensor_index_by_tuple(data, args)
-    if len(args) > 1:
-        const_utils.raise_value_error("Incorrect number of indices for array")
-    output = _tensor_index_by_integer(F.reshape(data, (-1,)), args[0])
-    return TensorToScalar()(output)
-
+    if data.shape == (1,):
+        return TensorToScalar()(data[0])
+    exp_msg = const_utils.gen_exception_msg("The tensor should have only one element. "
+                                            "But the shape of input tensor is {}.", data.shape)
+    const_utils.raise_value_error(exp_msg)
+    return None
 
 def tensor_itemset(data, *args):
     """Tensor setitem by index and value."""

mindspore/ops/composite/multitype_ops/pow_impl.py

@@ -16,7 +16,6 @@
 """Implementation for internal polymorphism `pow` operations."""
 
 from __future__ import absolute_import
-from mindspore.ops.composite.multitype_ops import _compile_utils as utils
 from mindspore.ops.composite import base
 from mindspore.ops import functional as F
 
@@ -53,34 +52,6 @@ def _scalar_pow_tensor(x, y):
     return F.tensor_pow(x, y)
 
 
-@pow_.register("Tuple", "Tensor")
-def _tuple_pow_tensor(x, y):
-    """Returns x ** y where x is a tuple and y is a tensor. """
-    x = utils.sequence_to_tensor(x, y.dtype)
-    return F.tensor_pow(x, y)
-
-
-@pow_.register("Tensor", "Tuple")
-def _tensor_pow_tuple(x, y):
-    """Returns x ** y where x is a tensor and y is a tuple. """
-    y = utils.sequence_to_tensor(y, x.dtype)
-    return F.tensor_pow(x, y)
-
-
-@pow_.register("List", "Tensor")
-def _list_pow_tensor(x, y):
-    """Returns x ** y where x is a list and y is a tensor. """
-    x = utils.sequence_to_tensor(x, y.dtype)
-    return F.tensor_pow(x, y)
-
-
-@pow_.register("Tensor", "List")
-def _tensor_pow_list(x, y):
-    """Returns x ** y where x is a tensor and y is a list. """
-    y = utils.sequence_to_tensor(y, x.dtype)
-    return F.tensor_pow(x, y)
-
-
 @pow_.register_default()
 def default_pow(x, y):
     """Default function for pow."""

mindspore/ops/function/__init__.py

@@ -167,6 +167,9 @@ from .array_func import (
     top_k,
     deepcopy,
     arange_ext,
+    view_as,
+    type_as,
+    expand_as,
 )
 from .parameter_func import (
     assign,
@@ -317,6 +320,7 @@ from .math_func import (
     sinh,
     cosh,
     tanh,
+    tanh_,
     tanhshrink,
     asinh,
     arcsinh,
@@ -379,6 +383,7 @@ from .math_func import (
     atleast_1d,
     dstack,
     diff,
+    diff_ext,
     atleast_2d,
     cartesian_prod,
     atleast_3d,
@@ -400,6 +405,7 @@ from .math_func import (
     remainder,
     remainder_ext,
     iou,
+    rotated_iou,
     bmm,
     trapz,
     cholesky,
@@ -443,6 +449,7 @@ from .math_func import (
     tensor_dot,
     vecdot,
     dot,
+    isnan_ext,
     batch_dot,
     eps,
 )
@@ -460,6 +467,7 @@ from .nn_func import (
     max_pool2d,
     max_pool3d,
     batch_norm,
+    add_rms_norm,
     rms_norm,
     bidense,
     celu,
@@ -485,6 +493,8 @@ from .nn_func import (
     soft_shrink,
     is_floating_point,
     incre_flash_attention,
+    prompt_flash_attention,
+    flash_attention_score,
     intopk,
     interpolate,
     upsample,
@@ -531,6 +541,7 @@ from .nn_func import (
     sigmoid,
     logsigmoid,
     relu,
+    relu_,
     relu6,
     rrelu,
     conv3d,
@@ -723,6 +734,7 @@ from .clip_func import (
     clip_by_value,
     clip_by_norm,
     clamp,
+    clamp_,
     clip,
     clip_by_global_norm,
 )