mindspore 2.3.0rc1__cp38-none-any.whl → 2.3.0rc2__cp38-none-any.whl

mindspore/ops/auto_generate/pyboost_inner_prim.py

@@ -17,10 +17,36 @@ from mindspore.common._stub_tensor import _convert_stub
 from mindspore.ops.auto_generate.gen_arg_handler import *
 from mindspore._c_expression import ArgMaxWithValuePrim_
 from mindspore._c_expression import ArgMinWithValuePrim_
+from mindspore._c_expression import BatchMatMulPrim_
+from mindspore._c_expression import BatchNormGradExtPrim_
 from mindspore._c_expression import BroadcastToPrim_
 from mindspore._c_expression import ConcatPrim_
+from mindspore._c_expression import ConvolutionGradPrim_
+from mindspore._c_expression import ConvolutionPrim_
+from mindspore._c_expression import FFNExtPrim_
+from mindspore._c_expression import FlashAttentionScoreGradPrim_
+from mindspore._c_expression import FlashAttentionScorePrim_
+from mindspore._c_expression import GridSampler2DGradPrim_
+from mindspore._c_expression import GridSampler2DPrim_
+from mindspore._c_expression import GridSampler3DGradPrim_
+from mindspore._c_expression import GridSampler3DPrim_
+from mindspore._c_expression import MatMulPrim_
+from mindspore._c_expression import MaxPoolGradWithIndicesPrim_
+from mindspore._c_expression import MaxPoolGradWithMaskPrim_
+from mindspore._c_expression import MaxPoolWithIndicesPrim_
+from mindspore._c_expression import MaxPoolWithMaskPrim_
+from mindspore._c_expression import OneHotExtPrim_
+from mindspore._c_expression import QuantBatchMatmulPrim_
+from mindspore._c_expression import ReduceAllPrim_
 from mindspore._c_expression import ReduceAnyPrim_
+from mindspore._c_expression import ReverseV2Prim_
 from mindspore._c_expression import SoftmaxPrim_
+from mindspore._c_expression import StackExtPrim_
+from mindspore._c_expression import TrilPrim_
+from mindspore._c_expression import TriuPrim_
+from mindspore._c_expression import UpsampleTrilinear3DGradPrim_
+from mindspore._c_expression import UpsampleTrilinear3DPrim_
+from mindspore._c_expression import WeightQuantBatchMatmulPrim_
 
 
 class _PyboostArgMaxWithValuePrim(ArgMaxWithValuePrim_):
@@ -41,6 +67,24 @@ class _PyboostArgMinWithValuePrim(ArgMinWithValuePrim_):
 argmin_with_value_impl = _PyboostArgMinWithValuePrim()
 
 
+class _PyboostBatchMatMulPrim(BatchMatMulPrim_):
+    def __call__(self, x, y, transpose_a, transpose_b):
+
+        return _convert_stub(super().__call__(x, y, transpose_a, transpose_b))
+
+
+batch_mat_mul_impl = _PyboostBatchMatMulPrim()
+
+
+class _PyboostBatchNormGradExtPrim(BatchNormGradExtPrim_):
+    def __call__(self, dout, input, weight, running_mean, running_var, saved_mean, saved_rstd, training, eps):
+
+        return _convert_stub(super().__call__(dout, input, weight, running_mean, running_var, saved_mean, saved_rstd, training, eps))
+
+
+batch_norm_grad_ext_impl = _PyboostBatchNormGradExtPrim()
+
+
 class _PyboostBroadcastToPrim(BroadcastToPrim_):
     def __call__(self, input, shape):
 
@@ -59,6 +103,181 @@ class _PyboostConcatPrim(ConcatPrim_):
 concat_impl = _PyboostConcatPrim()
 
 
+class _PyboostConvolutionGradPrim(ConvolutionGradPrim_):
+    def __call__(self, dout, input, weight, bias, stride, padding, dilation, transposed, output_padding, groups, output_mask):
+        converted_stride = to_strides(stride)
+        converted_padding = to_2d_paddings(padding)
+        converted_dilation = to_dilations(dilation)
+        converted_output_padding = to_output_padding(output_padding)
+        return _convert_stub(super().__call__(dout, input, weight, bias, stride, padding, dilation, transposed, output_padding, groups, output_mask))
+
+
+convolution_grad_impl = _PyboostConvolutionGradPrim()
+
+
+class _PyboostConvolutionPrim(ConvolutionPrim_):
+    def __call__(self, input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
+        converted_stride = to_strides(stride)
+        converted_padding = to_2d_paddings(padding)
+        converted_dilation = to_dilations(dilation)
+        converted_output_padding = to_output_padding(output_padding)
+        return _convert_stub(super().__call__(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups))
+
+
+convolution_impl = _PyboostConvolutionPrim()
+
+
+class _PyboostFFNExtPrim(FFNExtPrim_):
+    def __call__(self, x, weight1, weight2, expertTokens, bias1, bias2, scale, offset, deqScale1, deqScale2, antiquant_scale1, antiquant_scale2, antiquant_offset1, antiquant_offset2, activation, inner_precise):
+        converted_activation = str_to_enum(activation)
+        return _convert_stub(super().__call__(x, weight1, weight2, expertTokens, bias1, bias2, scale, offset, deqScale1, deqScale2, antiquant_scale1, antiquant_scale2, antiquant_offset1, antiquant_offset2, activation, inner_precise))
+
+
+ffn_ext_impl = _PyboostFFNExtPrim()
+
+
+class _PyboostFlashAttentionScoreGradPrim(FlashAttentionScoreGradPrim_):
+    def __call__(self, query, key, value, dy, pse_shift, drop_mask, padding_mask, atten_mask, softmax_max, softmax_sum, softmax_in, attention_in, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode):
+        converted_input_layout = str_to_enum(input_layout)
+        return _convert_stub(super().__call__(query, key, value, dy, pse_shift, drop_mask, padding_mask, atten_mask, softmax_max, softmax_sum, softmax_in, attention_in, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode))
+
+
+flash_attention_score_grad_impl = _PyboostFlashAttentionScoreGradPrim()
+
+
+class _PyboostFlashAttentionScorePrim(FlashAttentionScorePrim_):
+    def __call__(self, query, key, value, real_shift, drop_mask, padding_mask, attn_mask, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode):
+        converted_input_layout = str_to_enum(input_layout)
+        return _convert_stub(super().__call__(query, key, value, real_shift, drop_mask, padding_mask, attn_mask, prefix, actual_seq_qlen, actual_seq_kvlen, head_num, keep_prob, scale_value, pre_tokens, next_tokens, inner_precise, input_layout, sparse_mode))
+
+
+flash_attention_score_impl = _PyboostFlashAttentionScorePrim()
+
+
+class _PyboostGridSampler2DGradPrim(GridSampler2DGradPrim_):
+    def __call__(self, grad, input_x, grid, interpolation_mode, padding_mode, align_corners):
+        converted_interpolation_mode = str_to_enum(interpolation_mode)
+        converted_padding_mode = str_to_enum(padding_mode)
+        return _convert_stub(super().__call__(grad, input_x, grid, interpolation_mode, padding_mode, align_corners))
+
+
+grid_sampler_2d_grad_impl = _PyboostGridSampler2DGradPrim()
+
+
+class _PyboostGridSampler2DPrim(GridSampler2DPrim_):
+    def __call__(self, input_x, grid, interpolation_mode, padding_mode, align_corners):
+        converted_interpolation_mode = str_to_enum(interpolation_mode)
+        converted_padding_mode = str_to_enum(padding_mode)
+        return _convert_stub(super().__call__(input_x, grid, interpolation_mode, padding_mode, align_corners))
+
+
+grid_sampler_2d_impl = _PyboostGridSampler2DPrim()
+
+
+class _PyboostGridSampler3DGradPrim(GridSampler3DGradPrim_):
+    def __call__(self, grad, input_x, grid, interpolation_mode, padding_mode, align_corners):
+        converted_interpolation_mode = str_to_enum(interpolation_mode)
+        converted_padding_mode = str_to_enum(padding_mode)
+        return _convert_stub(super().__call__(grad, input_x, grid, interpolation_mode, padding_mode, align_corners))
+
+
+grid_sampler_3d_grad_impl = _PyboostGridSampler3DGradPrim()
+
+
+class _PyboostGridSampler3DPrim(GridSampler3DPrim_):
+    def __call__(self, input_x, grid, interpolation_mode, padding_mode, align_corners):
+        converted_interpolation_mode = str_to_enum(interpolation_mode)
+        converted_padding_mode = str_to_enum(padding_mode)
+        return _convert_stub(super().__call__(input_x, grid, interpolation_mode, padding_mode, align_corners))
+
+
+grid_sampler_3d_impl = _PyboostGridSampler3DPrim()
+
+
+class _PyboostMatMulPrim(MatMulPrim_):
+    def __call__(self, input, mat2, transpose_a, transpose_b):
+
+        return _convert_stub(super().__call__(input, mat2, transpose_a, transpose_b))
+
+
+matmul_impl = _PyboostMatMulPrim()
+
+
+class _PyboostMaxPoolGradWithIndicesPrim(MaxPoolGradWithIndicesPrim_):
+    def __call__(self, x, grad, argmax, kernel_size, strides, pads, dilation, ceil_mode, argmax_type):
+        converted_kernel_size = to_kernel_size(kernel_size)
+        converted_strides = to_strides(strides)
+        converted_pads = to_output_padding(pads)
+        converted_dilation = to_dilations(dilation)
+        return _convert_stub(super().__call__(x, grad, argmax, kernel_size, strides, pads, dilation, ceil_mode, argmax_type))
+
+
+max_pool_grad_with_indices_impl = _PyboostMaxPoolGradWithIndicesPrim()
+
+
+class _PyboostMaxPoolGradWithMaskPrim(MaxPoolGradWithMaskPrim_):
+    def __call__(self, x, grad, mask, kernel_size, strides, pads, dilation, ceil_mode, argmax_type):
+        converted_kernel_size = to_kernel_size(kernel_size)
+        converted_strides = to_strides(strides)
+        converted_pads = to_output_padding(pads)
+        converted_dilation = to_dilations(dilation)
+        return _convert_stub(super().__call__(x, grad, mask, kernel_size, strides, pads, dilation, ceil_mode, argmax_type))
+
+
+max_pool_grad_with_mask_impl = _PyboostMaxPoolGradWithMaskPrim()
+
+
+class _PyboostMaxPoolWithIndicesPrim(MaxPoolWithIndicesPrim_):
+    def __call__(self, x, kernel_size, strides, pads, dilation, ceil_mode, argmax_type):
+        converted_kernel_size = to_kernel_size(kernel_size)
+        converted_strides = to_strides(strides)
+        converted_pads = to_output_padding(pads)
+        converted_dilation = to_dilations(dilation)
+        return _convert_stub(super().__call__(x, kernel_size, strides, pads, dilation, ceil_mode, argmax_type))
+
+
+max_pool_with_indices_impl = _PyboostMaxPoolWithIndicesPrim()
+
+
+class _PyboostMaxPoolWithMaskPrim(MaxPoolWithMaskPrim_):
+    def __call__(self, x, kernel_size, strides, pads, dilation, ceil_mode, argmax_type):
+        converted_kernel_size = to_kernel_size(kernel_size)
+        converted_strides = to_strides(strides)
+        converted_pads = to_output_padding(pads)
+        converted_dilation = to_dilations(dilation)
+        return _convert_stub(super().__call__(x, kernel_size, strides, pads, dilation, ceil_mode, argmax_type))
+
+
+max_pool_with_mask_impl = _PyboostMaxPoolWithMaskPrim()
+
+
+class _PyboostOneHotExtPrim(OneHotExtPrim_):
+    def __call__(self, tensor, num_classes, on_value, off_value, axis):
+
+        return _convert_stub(super().__call__(tensor, num_classes, on_value, off_value, axis))
+
+
+one_hot_ext_impl = _PyboostOneHotExtPrim()
+
+
+class _PyboostQuantBatchMatmulPrim(QuantBatchMatmulPrim_):
+    def __call__(self, x1, x2, scale, offset, bias, transpose_x1, transpose_x2, dtype):
+
+        return _convert_stub(super().__call__(x1, x2, scale, offset, bias, transpose_x1, transpose_x2, dtype))
+
+
+quant_batch_matmul_impl = _PyboostQuantBatchMatmulPrim()
+
+
+class _PyboostReduceAllPrim(ReduceAllPrim_):
+    def __call__(self, input, axis, keep_dims):
+
+        return _convert_stub(super().__call__(input, axis, keep_dims))
+
+
+reduce_all_impl = _PyboostReduceAllPrim()
+
+
 class _PyboostReduceAnyPrim(ReduceAnyPrim_):
     def __call__(self, x, axis, keep_dims):
 
@@ -68,6 +287,15 @@ class _PyboostReduceAnyPrim(ReduceAnyPrim_):
 reduce_any_impl = _PyboostReduceAnyPrim()
 
 
+class _PyboostReverseV2Prim(ReverseV2Prim_):
+    def __call__(self, input, axis):
+
+        return _convert_stub(super().__call__(input, axis))
+
+
+reverse_v2_impl = _PyboostReverseV2Prim()
+
+
 class _PyboostSoftmaxPrim(SoftmaxPrim_):
     def __call__(self, input, axis):
 
@@ -75,3 +303,57 @@ class _PyboostSoftmaxPrim(SoftmaxPrim_):
 
 
 softmax_impl = _PyboostSoftmaxPrim()
+
+
+class _PyboostStackExtPrim(StackExtPrim_):
+    def __call__(self, tensors, dim):
+
+        return _convert_stub(super().__call__(tensors, dim))
+
+
+stack_ext_impl = _PyboostStackExtPrim()
+
+
+class _PyboostTrilPrim(TrilPrim_):
+    def __call__(self, input, diagonal):
+
+        return _convert_stub(super().__call__(input, diagonal))
+
+
+tril_impl = _PyboostTrilPrim()
+
+
+class _PyboostTriuPrim(TriuPrim_):
+    def __call__(self, input, diagonal):
+
+        return _convert_stub(super().__call__(input, diagonal))
+
+
+triu_impl = _PyboostTriuPrim()
+
+
+class _PyboostUpsampleTrilinear3DGradPrim(UpsampleTrilinear3DGradPrim_):
+    def __call__(self, dy, input_size, output_size, scales, align_corners):
+
+        return _convert_stub(super().__call__(dy, input_size, output_size, scales, align_corners))
+
+
+upsample_trilinear3d_grad_impl = _PyboostUpsampleTrilinear3DGradPrim()
+
+
+class _PyboostUpsampleTrilinear3DPrim(UpsampleTrilinear3DPrim_):
+    def __call__(self, x, output_size, scales, align_corners):
+
+        return _convert_stub(super().__call__(x, output_size, scales, align_corners))
+
+
+upsample_trilinear3d_impl = _PyboostUpsampleTrilinear3DPrim()
+
+
+class _PyboostWeightQuantBatchMatmulPrim(WeightQuantBatchMatmulPrim_):
+    def __call__(self, x, weight, antiquant_scale, antiquant_offset, quant_scale, quant_offset, bias, transpose_x, transpose_weight, antiquant_group_size):
+
+        return _convert_stub(super().__call__(x, weight, antiquant_scale, antiquant_offset, quant_scale, quant_offset, bias, transpose_x, transpose_weight, antiquant_group_size))
+
+
+weight_quant_batch_matmul_impl = _PyboostWeightQuantBatchMatmulPrim()

mindspore/ops/composite/multitype_ops/_compile_utils.py

@@ -137,7 +137,8 @@ def data_update_by_ops(transfer_type, arg, data, new_index, origin_data, value=N
     elif transfer_type == ValueTransferType.kGatherND:
         if isinstance(new_index, list):
             new_index = handle_multi_dim_index_tensor(new_index, arg)
-        data = F.gather_nd(data, Tensor(new_index))
+        new_index = format_index_tensor(new_index, (None, F.shape(data)[:F.shape(new_index)[-1]]))
+        data = F.gather_nd(data, new_index)
     elif transfer_type == ValueTransferType.kTensorScatterUpdate:
         if isinstance(new_index, list):
             new_index = handle_multi_dim_index_tensor(new_index, arg)

mindspore/ops/composite/multitype_ops/_constexpr_utils.py

@@ -459,13 +459,22 @@ def tuple_index_type_cnt(types, op_name):
 def check_value_elements(types):
     """Judges the type of all elements of the tuple."""
     tensor_number = 0
+    last_type = None
+    mix_but_no_tensor = False
     for ele in types:
         if isinstance(ele, mstype.TensorType):
             tensor_number += 1
         elif isinstance(ele, (list, tuple)):
             return MIXED
 
+        if last_type is None:
+            last_type = type(ele)
+        elif not isinstance(ele, last_type):
+            mix_but_no_tensor = True
+
     if tensor_number == 0:
+        if mix_but_no_tensor:
+            return MIXED
         return NO_TENSOR
     if tensor_number == len(types):
         return ALL_TENSOR

mindspore/ops/extend/__init__.py

@@ -33,13 +33,21 @@ from . import (
     nn_func,
 )
 
-from .array_func import gather, max, min
+from .array_func import gather, max, min, one_hot, narrow
 from .math_func import (
     baddbmm,
+    bmm,
     add,
     sub
 )
 
+from .nn_func import (
+    conv2d,
+    max_pool2d,
+    leaky_relu_ext,
+    batch_norm
+)
+
 __all__ = []
 __all__.extend(array_func.__all__)
 __all__.extend(math_func.__all__)

mindspore/ops/extend/array_func.py

@@ -18,12 +18,54 @@
 Array Operators
 
 """
-
+from mindspore.common import Tensor
 from mindspore.ops.operations.array_ops import ArgMaxWithValue, ArgMinWithValue
 from mindspore.ops._primitive_cache import _get_cache_prim
-from mindspore.ops.auto_generate.gen_ops_prim import gather_d_op
+from mindspore.ops.auto_generate.gen_ops_prim import gather_d_op, slice_ext_op, OneHotExt
+from mindspore.ops.auto_generate.gen_ops_def import max_, min_
+from mindspore import _checkparam as validator
+
 
 # define Primitive global variables
+def narrow(input, dim, start, length):
+    """
+    Returns a narrowed tensor from input tensor, and
+    the dimension axis is input from start to start + length.
+
+    Args:
+        input (Tensor): the tensor to narrow.
+        dim (int): dimension  along which to narrow.
+        start (int): the starting dimension.
+        length (int): the distance to the ending dimension.
+
+    Returns:
+        Tensor.
+
+        - output (Tensors) - The narrowed tensor.
+
+    Raises:
+        TypeError: If the input is not a tensor or tuple or list of tensors.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import ops
+        >>> from mindspore import Tensor
+        >>> x = Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], mindspore.int32)
+        >>> output = ops.narrow(x, 0, 0, 2)
+        >>> print(output)
+        [[ 1 2 3]
+         [ 4 5 6]]
+        >>> output = ops.narrow(x, 1, 1, 2)
+        >>> print(output)
+        [[ 2 3]
+         [ 5 6]
+         [ 8 9]]
+    """
+    validator.check_value_type("input", input, Tensor, "narrow")
+    return slice_ext_op(input, dim, start, start+length, 1)
 
 
 def gather(input, dim, index):
@@ -75,25 +117,32 @@ def gather(input, dim, index):
     return gather_d_op(input, dim, index)
 
 
-def max(input, dim, keepdim=False):
+def max(input, dim=None, keepdim=False):
     """
-    Calculates the maximum value along with the given axis for the input tensor.
+    Calculates the maximum value along with the given dimension for the input tensor.
 
     Args:
         input (Tensor): The input tensor, can be any dimension. Complex tensor is not supported for now.
-        dim (int): The dimension to reduce.
-        keepdim (bool): Whether to reduce dimension, if true, the output will keep same dimension with the input,
-            the output will reduce dimension if false. Default: ``False`` .
+        dim (int, optional): The dimension to reduce. Default: ``None`` .
+        keepdim (bool, optional): Whether to reduce dimension, if true, the output will keep same dimension
+            with the input, the output will reduce dimension if false. Default: ``False`` .
 
     Returns:
-        tuple (Tensor), tuple of 2 tensors, containing the maximum value of the input tensor and the corresponding
-        index.
+        Tensor if `dim` is the default value ``None`` , the maximum value of input tensor, with the shape :math:`()` ,
+        and same dtype as `input`.
+
+        tuple (Tensor) if `dim` is not the default value ``None`` , tuple of 2 tensors, containing the maximum
+        value of the input tensor along the given dimension `dim` and the corresponding index:
 
-        - values (Tensor) - The maximum value of input tensor, with same dtype as `input`. If `keepdim`
-          is true, the shape of output tensors is :math:`(input_1, input_2, ..., input_{axis-1}, 1, input_{axis+1},
-          ..., input_N)` . Otherwise, the shape is :math:`(input_1, input_2, ..., input_{axis-1}, input_{axis+1},
-          ..., input_N)` .
-        - index (Tensor) - The index for the maximum value of the input tensor, with the same shape as `values`.
+        - **values (Tensor)** - The maximum value of input tensor along the given dimension `dim`, with same dtype as
+          `input`. If `keepdim` is ``True`` , the shape of output tensors is :math:`(input_1, input_2, ...,
+          input_{axis-1}, 1, input_{axis+1}, ..., input_N)` . Otherwise, the shape is :math:`(input_1, input_2, ...,
+          input_{axis-1}, input_{axis+1}, ..., input_N)` .
+        - **index (Tensor)** - The index for the maximum value of the input tensor along the given dimension `dim`, with
+          the same shape as `values`.
+
+    Raises:
+        ValueError: If `dim` is the default value ``None`` and `keepdim` is not ``False`` .
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -108,30 +157,41 @@ def max(input, dim, keepdim=False):
         >>> print(output, index)
         [[3.2 0.4 0.4 2.9 4. ]] [[1 1 0 1 1]]
     """
+    if dim is None:
+        if keepdim is not False:
+            raise ValueError(f"For 'max', the `keepdim` must be False when the `dim` is None, but got {keepdim}")
+        return max_(input)
     argmax_with_value_op = _get_cache_prim(ArgMaxWithValue)(dim, keepdim)
     indices, values = argmax_with_value_op(input)
     return values, indices
 
 
-def min(input, dim, keepdim=False):
+def min(input, dim=None, keepdim=False):
     """
-    Calculates the minimum value along with the given axis for the input tensor.
+    Calculates the minimum value along with the given dimension for the input tensor.
 
     Args:
         input (Tensor): The input tensor, can be any dimension. Complex tensor is not supported for now.
-        dim (int): The dimension to reduce.
-        keepdim (bool): Whether to reduce dimension, if true, the output will keep same dimension with the input,
-            the output will reduce dimension if false. Default: ``False`` .
+        dim (int, optional): The dimension to reduce. Default: ``None`` .
+        keepdim (bool, optional): Whether to reduce dimension, if true, the output will keep same dimension
+            with the input, the output will reduce dimension if false. Default: ``False`` .
 
     Returns:
-        tuple (Tensor), tuple of 2 tensors, containing the minimum value of the input tensor and the corresponding
-        index.
+        Tensor if `dim` is the default value ``None`` , the minimum value of input tensor, with the shape :math:`()` ,
+        and same dtype as `input`.
+
+        tuple (Tensor) if `dim` is not the default value ``None`` , tuple of 2 tensors, containing the minimum value
+        of the input tensor along the given dimension `dim` and the corresponding index:
 
-        - values (Tensor) - The minimum value of input tensor, with same dtype as `input`. If `keepdim`
-          is true, the shape of output tensors is :math:`(input_1, input_2, ..., input_{axis-1}, 1, input_{axis+1},
-          ..., input_N)` . Otherwise, the shape is :math:`(input_1, input_2, ..., input_{axis-1}, input_{axis+1},
-          ..., input_N)` .
-        - index (Tensor) - The index for the minimum value of the input tensor, with same shape as `values`.
+        - **values (Tensor)** - The minimum value of input tensor along the given dimension `dim`, with same dtype as
+          `input`. If `keepdim` is ``True`` , the shape of output tensors is :math:`(input_1, input_2, ...,
+          input_{axis-1}, 1, input_{axis+1}, ..., input_N)` . Otherwise, the shape is :math:`(input_1, input_2, ...,
+          input_{axis-1}, input_{axis+1}, ..., input_N)` .
+        - **index (Tensor)** - The index for the minimum value of the input tensor along the given dimension `dim`,
+          with the same shape as `values`.
+
+    Raises:
+        ValueError: If `dim` is the default value ``None`` and `keepdim` is not ``False`` .
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -145,8 +205,55 @@ def min(input, dim, keepdim=False):
         >>> print(output, index)
         [0.0] [0]
     """
+    if dim is None:
+        if keepdim is not False:
+            raise ValueError(f"For 'min', the `keepdim` must be False when the `dim` is None, but got {keepdim}")
+        return min_(input)
     argmin_with_value_op = _get_cache_prim(ArgMinWithValue)(dim, keepdim)
     indices, values = argmin_with_value_op(input)
     return values, indices
 
-__all__ = ['gather', 'max', 'min']
+
+def one_hot(tensor, num_classes):
+    r"""
+    Computes a one-hot tensor.
+
+    The locations represented by tensor in `tensor` take value `1`, while all
+    other locations take value `0`.
+
+    Args:
+        tensor (Tensor): A tensor of indices. Tensor of shape :math:`(X_0, \ldots, X_n)`.
+            Data type must be int32 or int64.
+        num_classes (Union[int, Tensor]): A scalar defining the depth of the one-hot dimension.
+
+    Returns:
+        Tensor, one-hot tensor.
+
+    Raises:
+        TypeError: If `num_classes` is not an int.
+        TypeError: If dtype of `tensor` is not int32 or int64.
+        ValueError: If `num_classes` is less than 0.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> import mindspore.ops as ops
+        >>> from mindspore import Tensor
+        >>> tensor = Tensor(np.array([0, 1, 2]), mindspore.int32)
+        >>> num_classes = 3
+        >>> output = ops.extend.one_hot(tensor, num_classes)
+        >>> print(output)
+        [[1. 0. 0.]
+        [0. 1. 0.]
+        [0. 0. 1.]]
+    """
+    on_value = Tensor(1, dtype=tensor.dtype)
+    off_value = Tensor(0, dtype=tensor.dtype)
+    onehot = _get_cache_prim(OneHotExt)(-1)
+    return onehot(tensor, num_classes, on_value, off_value)
+
+
+__all__ = ['gather', 'max', 'min', 'one_hot']

mindspore/ops/extend/math_func.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2023 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,7 +20,7 @@ Math Operators with better performance
 """
 
 from mindspore.ops import auto_generate as P
-from mindspore.ops.auto_generate.gen_ops_def import add_ext as add, sub_ext as sub
+from mindspore.ops.auto_generate.gen_ops_def import add_ext as add, sub_ext as sub, bmm_ext as bmm
 
 
 # define Primitive global variables
@@ -73,4 +73,4 @@ def baddbmm(input, batch1, batch2, beta=1, alpha=1):
     return P.baddbmm(input, batch1, batch2, beta, alpha)
 
 
-__all__ = ['baddbmm', 'add', 'sub']
+__all__ = ['baddbmm', 'add', 'sub', 'bmm']