mindspore 2.3.0__cp39-none-any.whl → 2.3.0rc2__cp39-none-any.whl

mindspore/ops/operations/array_ops.py

@@ -41,7 +41,7 @@ from ..auto_generate import (ExpandDims, Reshape, TensorShape, Transpose, Gather
                              GatherNd, GatherD, Range, MaskedFill, RightShift, NonZero,
                              ResizeNearestNeighbor, Identity, Split, CumSum, CumProd,
                              Cummax, Cummin, Argmin, Concat, UnsortedSegmentSum, ScalarToTensor,
-                             Triu, BroadcastTo, StridedSlice, Select, TopkExt, SearchSorted)
+                             Tril, Triu, BroadcastTo, StridedSlice, Select, TopkExt)
 from .manually_defined import Rank, Shape, Tile, Cast, Ones, Zeros
 from ..auto_generate import ArgMaxWithValue, ArgMinWithValue
 
@@ -1912,9 +1912,12 @@ class Stack(PrimitiveWithInfer):
         dtype = x_type[0]
         if tuple_value is not None and None not in tuple_value:
             for item in tuple_value:
-                npy_item = item.asnumpy()
+                npy_item = item.asnumpy() if item.dtype != mstype.bfloat16 else item.float().asnumpy()
                 input_array.append(npy_item)
-            infered_value = Tensor(np.stack(input_array, axis=self.axis))
+            if dtype == mstype.TensorType(mstype.bfloat16):
+                infered_value = Tensor(np.stack(input_array, axis=self.axis), mstype.bfloat16)
+            else:
+                infered_value = Tensor(np.stack(input_array, axis=self.axis))
 
         shape = all_shape.get('shape') if isinstance(all_shape, dict) else all_shape
         out = {'shape': shape,
@@ -4054,7 +4057,7 @@ class EditDistance(Primitive):
         >>> import numpy as np
         >>> from mindspore import Tensor
         >>> import mindspore.nn as nn
-        >>> from mindspore import ops
+        >>> import mindspore.ops as ops
         >>> class EditDistance(nn.Cell):
         ...     def __init__(self, hypothesis_shape, truth_shape, normalize=True):
         ...         super(EditDistance, self).__init__()
@@ -4403,6 +4406,60 @@ class MaskedSelect(PrimitiveWithCheck):
         validator.check_tensor_dtype_valid('x', x_dtype, (mstype.bool_,) + mstype.number_type, self.name)
 
 
+class SearchSorted(Primitive):
+    """
+    Returns the indices correspond to the positions where the given numbers in `values` should be inserted
+    into `sorted_sequence` so that the order of the sequence is maintained.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Refer to :func:`mindspore.ops.searchsorted` for more details.
+
+    Args:
+        dtype (:class:`mindspore.dtype`, optional): Output data type. An optional data type of
+            ``mstype.int32`` and ``mstype.int64``. Default: ``mstype.int64``.
+        right (bool, optional): Search Strategy. If ``True`` , return the last suitable index found;
+            if ``False`` , return the first such index. Default: ``False`` .
+
+    Inputs:
+        - **sorted_sequence** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R-1, x_R)` or `(x_1)`.
+          It must contain a monotonically increasing sequence on the innermost dimension.
+        - **values** (Tensor) - The value that should be inserted.
+          The shape of tensor is :math:`(x_1, x_2, ..., x_R-1, x_S)`.
+
+    Outputs:
+        Tensor containing the indices from the innermost dimension of `sorted_sequence` such that,
+        if insert the corresponding value in the `values` tensor, the order of `sorted_sequence` would be preserved,
+        whose datatype is int32 if out_int32 is True, otherwise int64, and shape is the same as the shape of `values`.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, ops
+        >>> sorted_sequence = Tensor(np.array([[0, 1, 3, 5, 7], [2, 4, 6, 8, 10]]), mindspore.float32)
+        >>> values = Tensor(np.array([[3, 6, 9], [3, 6, 9]]), mindspore.float32)
+        >>> output = ops.SearchSorted()(sorted_sequence, values)
+        >>> print(output)
+        [[2 4 5]
+         [1 2 4]]
+    """
+
+    @prim_attr_register
+    def __init__(self, dtype=mstype.int64, right=False):
+        """Initialize SearchSorted"""
+        validator.check_value_type("dtype", dtype, [mstype.Type], self.name)
+        valid_values = (mstype.int64, mstype.int32)
+        self.dtype = validator.check_type_name(
+            "dtype", dtype, valid_values, self.name)
+        validator.check_value_type('right', right, [bool], self.name)
+        self.init_prim_io_names(
+            inputs=['sorted_sequence', 'values'], outputs=['output'])
+
+
 class _TensorScatterOp(PrimitiveWithInfer):
     """
     Defines TensorScatter Base Operators
@@ -4996,7 +5053,7 @@ class TensorScatterElements(Primitive):
 
     Examples:
         >>> import mindspore
-        >>> from mindspore import ops
+        >>> import mindspore.ops as ops
         >>> from mindspore import Tensor
         >>> op = ops.TensorScatterElements(0, "none")
         >>> data = Tensor(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), mindspore.float32)
@@ -5008,7 +5065,7 @@ class TensorScatterElements(Primitive):
          [ 0.0  5.0  0.0]
          [ 7.0  0.0  0.0]]
         >>> import mindspore as ms
-        >>> from mindspore import ops
+        >>> import mindspore.ops as ops
         >>> from mindspore import Tensor
         >>> op = ops.TensorScatterElements(1, "add")
         >>> data = Tensor(np.array([[1, 2, 3, 4, 5]]), mindspore.float32)
@@ -5215,7 +5272,7 @@ class LowerBound(Primitive):
         >>> import mindspore
         >>> import numpy as np
         >>> from mindspore import Tensor
-        >>> from mindspore import ops
+        >>> import mindspore.ops as ops
         >>> lowerbound = ops.LowerBound(out_type = mindspore.int32)
         >>> sorted_x = Tensor(np.arange(12).reshape(3, 4).astype(np.int8))
         >>> values = Tensor(np.array([[3], [4], [8]]).astype(np.int8))
@@ -5268,7 +5325,7 @@ class UpperBound(Primitive):
         >>> import mindspore
         >>> import numpy as np
         >>> from mindspore import Tensor
-        >>> from mindspore import ops
+        >>> import mindspore.ops as ops
         >>> upperbound = ops.UpperBound(out_type = mindspore.int32)
         >>> sorted_x = Tensor(np.arange(12).reshape(3, 4).astype(np.int8))
         >>> values = Tensor(np.array([[3], [6], [9]]).astype(np.int8))
@@ -5354,80 +5411,6 @@ class LogSpace(Primitive):
         self.init_prim_io_names(inputs=['start', 'end'], outputs=['y'])
 
 
-class Tril(Primitive):
-    """
-    Returns the lower triangular portion of the 2-D matrix or the set of matrices
-    in a batch. The remaining elements of the resulting Tensor are assigned a value of 0.
-    The lower triangular section of the matrix comprises of the
-    elements present on and below the main diagonal.
-
-    .. warning::
-        This is an experimental API that is subject to change or deletion.
-
-    Args:
-        diagonal (int, optional): An optional attribute indicates the diagonal to consider, default: ``0`` ,
-            indicating the main diagonal.
-
-    Inputs:
-        - **x** (Tensor) - The input tensor with shape :math:`(M, N, *)`
-          where :math:`*` means any number of additional dimensions.
-
-    Outputs:
-        Tensor, the same shape and data type as the input `x`.
-
-    Raises:
-        TypeError: If `x` is not a Tensor.
-        TypeError: If `diagonal` is not an int.
-        ValueError: If the rank of `x` is less than 2.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-
-    Examples:
-        >>> import numpy as np
-        >>> from mindspore import Tensor, ops
-        >>> x = Tensor(np.array([[ 1,  2,  3,  4],
-        ...                      [ 5,  6,  7,  8],
-        ...                      [10, 11, 12, 13],
-        ...                      [14, 15, 16, 17]]))
-        >>> tril = ops.Tril()
-        >>> result = tril(x)
-        >>> print(result)
-        [[ 1  0  0  0]
-         [ 5  6  0  0]
-         [10 11 12  0]
-         [14 15 16 17]]
-        >>> x = Tensor(np.array([[ 1,  2,  3,  4],
-        ...                      [ 5,  6,  7,  8],
-        ...                      [10, 11, 12, 13],
-        ...                      [14, 15, 16, 17]]))
-        >>> tril = ops.Tril(diagonal=1)
-        >>> result = tril(x)
-        >>> print(result)
-        [[ 1  2  0  0]
-         [ 5  6  7  0]
-         [10 11 12 13]
-         [14 15 16 17]]
-        >>> x = Tensor(np.array([[ 1,  2,  3,  4],
-        ...                      [ 5,  6,  7,  8],
-        ...                      [10, 11, 12, 13],
-        ...                      [14, 15, 16, 17]]))
-        >>> tril = ops.Tril(diagonal=-1)
-        >>> result = tril(x)
-        >>> print(result)
-        [[ 0  0  0  0]
-         [ 5  0  0  0]
-         [10 11  0  0]
-         [14 15 16  0]]
-    """
-
-    @prim_attr_register
-    def __init__(self, diagonal=0):
-        """Initialize Tril."""
-        self.init_prim_io_names(inputs=["x"], outputs=["y"])
-        validator.check_value_type("diagonal", diagonal, [int], self.name)
-
-
 class IndexFill(Primitive):
     """
     Fills the elements under the `dim` dimension of the input Tensor `x` with the input `value`