mindspore 2.2.11__cp39-cp39-win_amd64.whl → 2.3.0__cp39-cp39-win_amd64.whl

mindspore/ops/function/random_func.py

@@ -1,4 +1,4 @@
-# Copyright 2022 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.
@@ -25,9 +25,27 @@ from mindspore.ops.composite.multitype_ops import _constexpr_utils as const_util
 from mindspore.common import dtype as mstype
 from mindspore.common.seed import _get_graph_seed
 from mindspore.common.tensor import Tensor
-from mindspore.ops.operations.random_ops import RandomShuffle, RandomChoiceWithMask, RandpermV2
-from mindspore.ops._primitive_cache import _get_cache_prim
+from mindspore.ops.operations.random_ops import RandomShuffle, RandomChoiceWithMask
 from mindspore.common.api import _function_forbid_reuse
+from mindspore.ops.auto_generate import randperm
+from mindspore.common.generator import default_generator
+from mindspore.ops.auto_generate import UniformExt, NormalTensorTensor, \
+    NormalTensorFloat, NormalFloatTensor, NormalFloatFloat, RandExt, RandLikeExt
+
+normal_tensor_tensor_op = NormalTensorTensor()
+normal_tensor_float_op = NormalTensorFloat()
+normal_float_tensor_op = NormalFloatTensor()
+normal_float_float_op = NormalFloatFloat()
+cast_ = P.Cast()
+log_ = P.Log()
+real_div_ = P.RealDiv()
+reshape_ = P.Reshape()
+shape_ = P.Shape()
+top_k_ = P.TopK()
+uniform_ = UniformExt()
+rand_ext_ = RandExt()
+rand_like_ext_ = RandLikeExt()
+generator_step_ = Tensor(10, mstype.int64)
 
 
 @constexpr
@@ -77,7 +95,8 @@ def random_gamma(shape, alpha, seed=None):
     """
     seed1, seed2 = _get_seed(seed, "random_gamma")
     random_gamma_op = P.RandomGamma(seed1, seed2)
-    random_gamma_op = _set_prim_op_user_data(random_gamma_op, "random_cache", False)
+    random_gamma_op = _set_prim_op_user_data(
+        random_gamma_op, "random_cache", False)
     output = random_gamma_op(shape, alpha)
     return output
 
@@ -124,7 +143,8 @@ def standard_laplace(shape, seed=None):
     """
     seed1, seed2 = _get_seed(seed, "standard_laplace")
     standard_laplace_op = P.StandardLaplace(seed=seed1, seed2=seed2)
-    standard_laplace_op = _set_prim_op_user_data(standard_laplace_op, "random_cache", False)
+    standard_laplace_op = _set_prim_op_user_data(
+        standard_laplace_op, "random_cache", False)
     return standard_laplace_op(shape)
 
 
@@ -163,7 +183,8 @@ def random_categorical(logits, num_sample, seed=0, dtype=mstype.int64):
         (10, 8)
     """
     random_categorical_ = P.RandomCategorical(dtype)
-    random_categorical_ = _set_prim_op_user_data(random_categorical_, "random_cache", False)
+    random_categorical_ = _set_prim_op_user_data(
+        random_categorical_, "random_cache", False)
     return random_categorical_(logits, num_sample, seed)
 
 
@@ -223,10 +244,44 @@ def multinomial_with_replacement(x, seed, offset, numsamples, replacement=False)
         offset = Tensor(offset, dtype=mstype.int64)
     multinomial_with_replacement_ = P.MultinomialWithReplacement(numsamples=numsamples,
                                                                  replacement=replacement)
-    multinomial_with_replacement_ = _set_prim_op_user_data(multinomial_with_replacement_, "random_cache", False)
+    multinomial_with_replacement_ = _set_prim_op_user_data(
+        multinomial_with_replacement_, "random_cache", False)
     return multinomial_with_replacement_(x, seed, offset)
 
 
+@_function_forbid_reuse
+def uniform_ext(tensor, a, b, generator=None):
+    """
+    Generates random numbers in the half-open interval [a, b).
+
+    Args:
+        tensor (Tensor): The origin input tensor.
+        a (number): The lower bound of the interval.
+        b (number): The upper bound of the interval.
+        generator (Generator, optional): The random seed. Default: None.
+
+    Raises:
+        TypeError: If `a` is larger than `b`.
+
+    Returns:
+        Tensor, with the same shape as tensor.
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import ops
+        >>> x = ops.ones((4, 2))
+        >>> generator = mindspore.Generator()
+        >>> generator.manual_seed(100)
+        >>> result = ops.function.random_func.uniform_ext(x, 1., 2., generator)
+        >>> print(result.shape)
+        (4, 2)
+    """
+    if generator is None:
+        generator = default_generator
+    seed, offset = generator._step(generator_step_)  # pylint: disable=protected-access
+    return uniform_(tensor, a, b, seed, offset)
+
+
 @_function_forbid_reuse
 def uniform(shape, minval, maxval, seed=None, dtype=mstype.float32):
     """
@@ -284,21 +339,25 @@ def uniform(shape, minval, maxval, seed=None, dtype=mstype.float32):
         (3, 2, 2)
     """
     if not isinstance(minval, Tensor) or not isinstance(maxval, Tensor):
-        raise TypeError(f"For functional operator[uniform], the input[minval] and input[maxval] must be a Tensor.")
+        raise TypeError(
+            f"For functional operator[uniform], the input[minval] and input[maxval] must be a Tensor.")
 
     minval_dtype = F.dtype(minval)
     maxval_dtype = F.dtype(maxval)
-    const_utils.check_type_valid(dtype, [mstype.int32, mstype.float32], 'uniform')
+    const_utils.check_type_valid(
+        dtype, [mstype.int32, mstype.float32], 'uniform')
     const_utils.check_tensors_dtype_same(minval_dtype, dtype, "uniform")
     const_utils.check_tensors_dtype_same(maxval_dtype, dtype, "uniform")
     seed1, seed2 = _get_seed(seed, "uniform")
     if const_utils.is_same_type(dtype, mstype.int32):
         random_uniform = P.UniformInt(seed1, seed2)
-        random_uniform = _set_prim_op_user_data(random_uniform, "random_cache", False)
+        random_uniform = _set_prim_op_user_data(
+            random_uniform, "random_cache", False)
         value = random_uniform(shape, minval, maxval)
     else:
         uniform_real = P.UniformReal(seed1, seed2)
-        uniform_real = _set_prim_op_user_data(uniform_real, "random_cache", False)
+        uniform_real = _set_prim_op_user_data(
+            uniform_real, "random_cache", False)
         uniform_real = uniform_real(shape)
         value = uniform_real * (maxval - minval) + minval
     return value
@@ -341,7 +400,8 @@ def standard_normal(shape, seed=None):
     """
     seed1, seed2 = _get_seed(seed, "standard_normal")
     standard_normal_op = P.StandardNormal(seed=seed1, seed2=seed2)
-    standard_normal_op = _set_prim_op_user_data(standard_normal_op, "random_cache", False)
+    standard_normal_op = _set_prim_op_user_data(
+        standard_normal_op, "random_cache", False)
     return standard_normal_op(shape)
 
 
@@ -374,11 +434,11 @@ def uniform_candidate_sampler(true_classes,
 
     Returns:
         - **sampled_candidates** (Tensor) - The sampled_candidates is independent of the true classes.
-          Shape: :math:`(num\_sampled, )` .
+          shape: :math:`(num\_sampled, )` .
         - **true_expected_count** (Tensor) - The expected counts under the sampling distribution of each
-          of true_classes. Shape: :math:`(batch\_size, num\_true)` .
+          of true_classes. shape: :math:`(batch\_size, num\_true)` .
         - **sampled_expected_count** (Tensor) - The expected counts under the sampling distribution of
-          each of sampled_candidates. Shape: :math:`(num\_sampled, )` .
+          each of sampled_candidates. shape: :math:`(num\_sampled, )` .
 
     Raises:
         TypeError: If neither `num_true` nor `num_sampled` is an int.
@@ -408,7 +468,8 @@ def uniform_candidate_sampler(true_classes,
                                            seed=seed,
                                            remove_accidental_hits=remove_accidental_hits)
     sampler_op = _set_prim_op_user_data(sampler_op, "random_cache", False)
-    sampled_candidates, true_expected_count, sampled_expected_count = sampler_op(true_classes)
+    sampled_candidates, true_expected_count, sampled_expected_count = sampler_op(
+        true_classes)
     return sampled_candidates, true_expected_count, sampled_expected_count
 
 
@@ -472,7 +533,8 @@ def random_poisson(shape, rate, seed=None, dtype=mstype.float32):
     """
     seed1, seed2 = _get_seed(seed, "random_poisson")
     prim_random_poisson = P.RandomPoisson(seed1, seed2, dtype)
-    prim_random_poisson = _set_prim_op_user_data(prim_random_poisson, "random_cache", False)
+    prim_random_poisson = _set_prim_op_user_data(
+        prim_random_poisson, "random_cache", False)
     value = prim_random_poisson(shape, rate)
     return value
 
@@ -507,7 +569,8 @@ def shuffle(x, seed=None):
     """
     seed, seed2 = _get_seed(seed, "shuffle")
     random_shuffle_ = RandomShuffle(seed=seed, seed2=seed2)
-    random_shuffle_ = _set_prim_op_user_data(random_shuffle_, "random_cache", False)
+    random_shuffle_ = _set_prim_op_user_data(
+        random_shuffle_, "random_cache", False)
     output = random_shuffle_(x)
     return output
 
@@ -561,7 +624,8 @@ def log_uniform_candidate_sampler(true_classes, num_true=1, num_sampled=5, uniqu
 
     """
 
-    sampler = P.LogUniformCandidateSampler(num_true, num_sampled, unique, range_max, seed)
+    sampler = P.LogUniformCandidateSampler(
+        num_true, num_sampled, unique, range_max, seed)
     sampler = _set_prim_op_user_data(sampler, "random_cache", False)
     return sampler(true_classes)
 
@@ -600,7 +664,7 @@ def choice_with_mask(input_x, count=256, seed=None):
     Examples:
         >>> import numpy as np
         >>> from mindspore import Tensor, ops
-        >>> input_x = Tensor(np.ones(shape=[240000, 4]).astype(np.bool))
+        >>> input_x = Tensor(np.ones(shape=[240000, 4]).astype(np.bool_))
         >>> output_y, output_mask = ops.choice_with_mask(input_x)
         >>> result = output_y.shape
         >>> print(result)
@@ -610,8 +674,10 @@ def choice_with_mask(input_x, count=256, seed=None):
         (256,)
     """
     seed1, seed2 = _get_seed(seed, "choice_with_mask")
-    choice_with_mask_ = RandomChoiceWithMask(count=count, seed=seed1, seed2=seed2)
-    choice_with_mask_ = _set_prim_op_user_data(choice_with_mask_, "random_cache", False)
+    choice_with_mask_ = RandomChoiceWithMask(
+        count=count, seed=seed1, seed2=seed2)
+    choice_with_mask_ = _set_prim_op_user_data(
+        choice_with_mask_, "random_cache", False)
     output = choice_with_mask_(input_x)
     return output
 
@@ -622,56 +688,54 @@ def is_cpu_backend():
     return context.get_context('device_target') == 'CPU'
 
 
-@_function_forbid_reuse
-def randperm(n, seed=0, offset=0, dtype=mstype.int64):
+def normal_ext(mean=0.0, std=1.0, size=None, generator=None):
     r"""
-    Generates random permutation of integers from 0 to n-1.
-
-    Returns the tensor with the determined shape inferred by n, the random numbers in it drawn from the data range
-    that a given type can represent.
-
-    .. warning::
-        This is an experimental API that is subject to change or deletion.
+    Generates random numbers according to the standard Normal (or Gaussian) random number distribution.
 
     Args:
-        n (Union[Tensor, int]): The input n Tensor with shape: () or (1,) and with data type of int64.
-            The value of `n` must be greater than zero.
-        seed (int, optional): Random seed. Default: ``0`` . When seed is -1(only negative value), offset is 0,
-            it's determined by time.
-        offset (int, optional): Offset to generate random numbers. Priority is higher than random seed.
-            Default: ``0`` . It must be non-negative.
-        dtype (mindspore.dtype, optional): The type of output.
-            Its value must be one of the following types: int32, int16, int8,
-            uint8, int64, float64, float32, float16. Default: mstype.int64.
+        mean (Union[float, Tensor], optional): Mean value of each element, the shape of the 'mean' tensor
+            should be the same as that of the 'std' tensor. Default: ``0.0``.
+        std (Union[float, Tensor], optional): Standard deviation for each element, the shape of the 'std' tensor
+            should be the same as that of the 'mean' tensor. The value of std should be greater than or equal to 0.
+            Default: ``1.0``.
+        size (tuple, optional): output size, where 'mean' and 'std' are constants. Default: ``None``.
+        generator (generator, optional): MindSpore generator. Default: ``None``.
 
     Returns:
-        Tensor. Its shape is specified by the required args `n`. Its type is specified by `dtype`.
-        Otherwise is default.
+        Outputs a tensor with the same shape as 'mean',
+        or when 'mean' and 'std' are constants and shape is specified as 'size'.
 
     Raises:
-        TypeError: If `dtype` is not allowed.
-        ValueError: If `n` is a negative or 0 element.
-        ValueError: If `seed` is a negative element.
-        ValueError: If `n` is larger than the maximal data of the set dtype.
+        TypeError: If `mean` or `std` is not Union[float, Tensor].
 
     Supported Platforms:
-        ``CPU``
+        ``Ascend``
 
     Examples:
+        >>> import mindspore
+        >>> import numpy as np
         >>> from mindspore import ops
-        >>> from mindspore import dtype as mstype
-        >>> n = 4
-        >>> seed = 0
-        >>> offset = 0
-        >>> output = ops.randperm(n, seed, offset, dtype=mstype.int64)
-        >>> print(output)
-        [1 0 2 3]
+        >>> from mindspore import Tensor
+        >>> mean = Tensor(np.array([1.0, 2.0, 3.0]), mindspore.float32)
+        >>> std = Tensor(np.array([1.0, 2.0, 3.0]), mindspore.float32)
+        >>> output = ops.function.random_func.normal_ext(mean, std)
+        >>> print(output.shape)
+        (3,)
     """
-    if not isinstance(n, Tensor):
-        n = Tensor(n)
-    randperm_ = RandpermV2(dtype=dtype)
-    randperm_ = _set_prim_op_user_data(randperm_, "random_cache", False)
-    return randperm_(n, seed, offset)
+    if generator is None:
+        generator = default_generator
+    seed, offset = generator._step(generator_step_)  # pylint: disable=protected-access
+
+    is_mean_tensor = isinstance(mean, Tensor)
+    is_std_tensor = isinstance(std, Tensor)
+
+    if is_mean_tensor and is_std_tensor:
+        return normal_tensor_tensor_op(mean, std, seed, offset)
+    if is_mean_tensor and not is_std_tensor:
+        return normal_tensor_float_op(mean, std, seed, offset)
+    if not is_mean_tensor and is_std_tensor:
+        return normal_float_tensor_op(mean, std, seed, offset)
+    return normal_float_float_op(mean, std, size, seed, offset)
 
 
 @_function_forbid_reuse
@@ -776,7 +840,8 @@ def laplace(shape, mean, lambda_param, seed=None):
     mean_dtype = F.dtype(mean)
     lambda_param_dtype = F.dtype(lambda_param)
     const_utils.check_tensors_dtype_same(mean_dtype, mstype.float32, "laplace")
-    const_utils.check_tensors_dtype_same(lambda_param_dtype, mstype.float32, "laplace")
+    const_utils.check_tensors_dtype_same(
+        lambda_param_dtype, mstype.float32, "laplace")
     seed1, seed2 = _get_seed(seed, "laplace")
     stdlaplace = P.StandardLaplace(seed1, seed2)
     stdlaplace = _set_prim_op_user_data(stdlaplace, "random_cache", False)
@@ -795,8 +860,8 @@ def gamma(shape, alpha, beta, seed=None):
         shape (tuple): The shape of random tensor to be generated.
         alpha (Tensor): The :math:`\alpha` distribution parameter. It should be greater than 0 with float32 data type.
         beta (Tensor): The :math:`\beta` distribution parameter. It should be greater than 0 with float32 data type.
-        seed (int): Seed is used as entropy source for the random number engines to generate
-          pseudo-random numbers, must be non-negative. Default: ``None`` , which will be treated as ``0`` .
+        seed (int, optional): Seed is used as entropy source for the random number engines to generate
+            pseudo-random numbers, must be non-negative. Default: ``None`` , which will be treated as ``0`` .
 
     Returns:
         Tensor. The shape should be equal to the broadcasted shape between the input `shape` and shapes
@@ -914,7 +979,7 @@ def rand(*size, dtype=None, seed=None):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> print(ops.rand((2,3)))
         [[4.1702199e-01 9.9718481e-01 7.2032452e-01]
          [9.3255734e-01 1.1438108e-04 1.2812445e-01]]
@@ -922,9 +987,9 @@ def rand(*size, dtype=None, seed=None):
     if dtype is None:
         dtype = mstype.float32
     elif dtype not in mstype.float_type:
-        raise ValueError(f"For 'rand', the 'dtype' must be a float type, but got {dtype}.")
+        raise ValueError(
+            f"For 'rand', the 'dtype' must be a float type, but got {dtype}.")
     shape = _generate_shapes(size)
-    cast_ = P.Cast()
     seed1, seed2 = _get_seed(seed, 'rand')
     rand_op = P.UniformReal(seed1, seed2)
     rand_op = _set_prim_op_user_data(rand_op, "random_cache", False)
@@ -966,13 +1031,14 @@ def rand_like(input, seed=None, *, dtype=None):
          [9.3255734e-01 1.1438108e-04 1.2812445e-01]]
     """
     if not isinstance(input, Tensor):
-        raise TypeError(f"For 'rand_like', the 'input' must be a Tensor, but got {type(input)}")
+        raise TypeError(
+            f"For 'rand_like', the 'input' must be a Tensor, but got {type(input)}")
     if dtype is None:
         dtype = input.dtype
     if dtype not in mstype.float_type:
-        raise ValueError(f"For 'rand_like', the 'dtype' must be a float type, but got {dtype}.")
+        raise ValueError(
+            f"For 'rand_like', the 'dtype' must be a float type, but got {dtype}.")
     shape = input.shape
-    cast_ = P.Cast()
     seed1, seed2 = _get_seed(seed, 'rand_like')
     rand_op = P.UniformReal(seed1, seed2)
     rand_op = _set_prim_op_user_data(rand_op, "random_cache", False)
@@ -980,6 +1046,76 @@ def rand_like(input, seed=None, *, dtype=None):
     return cast_(output, dtype)
 
 
+@_function_forbid_reuse
+def rand_ext(*size, generator=None, dtype=None):
+    r"""
+    Returns a new tensor that fills numbers from the uniform distribution over an interval :math:`[0, 1)`
+    based on the given shape and dtype.
+
+    Args:
+        size (Union[int, tuple(int), list(int)]): Shape of the new tensor, e.g. :math:`(2, 3)` or :math:`2`.
+
+    Keyword Args:
+        generator (:class:`mindspore.Generator`, optional): a pseudorandom number generator.
+            Default: ``None``, uses the default pseudorandom number generator.
+        dtype (:class:`mindspore.dtype`, optional): Designated tensor dtype, it must be float type. If None,
+            `mindspore.float32` will be applied. Default: ``None`` .
+
+    Returns:
+        Tensor, with the designated shape and dtype, filled with random numbers from the uniform distribution on
+        the interval :math:`[0, 1)`.
+
+    Raises:
+        ValueError: If `dtype` is not a `mstype.float_type` type.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore.ops as ops
+        >>> print(ops.function.random_func.rand_ext(2, 3).shape)
+        (2, 3)
+    """
+    if not generator:
+        generator = default_generator
+    seed, offset = generator._step(generator_step_)  # pylint: disable=protected-access
+    return rand_ext_(size, seed, offset, dtype)
+
+
+@_function_forbid_reuse
+def rand_like_ext(input, *, dtype=None):
+    r"""
+    Returns a new tensor that fills numbers from the uniform distribution over an interval :math:`[0, 1)`
+    based on the given dtype and shape of the input tensor.
+
+    Args:
+        input (Tensor): Input Tensor to specify the output shape and its default dtype.
+
+    Keyword Args:
+        dtype (:class:`mindspore.dtype`, optional): Designated tensor dtype, it must be float type. If None,
+            the same dtype of `input` will be applied. Default: ``None`` .
+
+    Returns:
+        Tensor, with the designated shape and dtype, filled with random numbers from the uniform distribution on
+        the interval :math:`[0, 1)`.
+
+    Raises:
+        ValueError: If `dtype` is not a `mstype.float_type` type.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore import Tensor, ops
+        >>> a = Tensor([[2, 3, 4], [1, 2, 3]])
+        >>> print(ops.function.random_func.rand_like_ext(a, dtype=ms.float32).shape)
+        (2, 3)
+    """
+    seed, offset = default_generator._step(generator_step_)  # pylint: disable=protected-access
+    return rand_like_ext_(input, seed, offset, dtype)
+
+
 @_function_forbid_reuse
 def randn(*size, dtype=None, seed=None):
     r"""
@@ -1007,7 +1143,7 @@ def randn(*size, dtype=None, seed=None):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> print(ops.randn((2, 2)))
         [[ 0.30639967 -0.42438635]
          [-0.4287376   1.3054721 ]]
@@ -1015,9 +1151,9 @@ def randn(*size, dtype=None, seed=None):
     if dtype is None:
         dtype = mstype.float32
     elif dtype not in mstype.float_type:
-        raise ValueError(f"For 'randn', the 'dtype' must be a float type, but got {dtype}.")
+        raise ValueError(
+            f"For 'randn', the 'dtype' must be a float type, but got {dtype}.")
     shape = _generate_shapes(size)
-    cast_ = P.Cast()
     seed1, seed2 = _get_seed(seed, 'randn')
     rand_op = P.StandardNormal(seed1, seed2)
     rand_op = _set_prim_op_user_data(rand_op, "random_cache", False)
@@ -1059,13 +1195,14 @@ def randn_like(input, seed=None, *, dtype=None):
          [-0.4287376   1.3054721   0.64747655]]
     """
     if not isinstance(input, Tensor):
-        raise TypeError(f"For 'randn_like', the 'input' must be a Tensor, but got {type(input)}")
+        raise TypeError(
+            f"For 'randn_like', the 'input' must be a Tensor, but got {type(input)}")
     if dtype is None:
         dtype = mstype.float32
     if dtype not in mstype.float_type:
-        raise ValueError(f"For 'randn_like', the 'dtype' must be a float type, but got {dtype}.")
+        raise ValueError(
+            f"For 'randn_like', the 'dtype' must be a float type, but got {dtype}.")
     shape = input.shape
-    cast_ = P.Cast()
     seed1, seed2 = _get_seed(seed, 'randn_like')
     rand_op = P.StandardNormal(seed1, seed2)
     rand_op = _set_prim_op_user_data(rand_op, "random_cache", False)
@@ -1103,7 +1240,7 @@ def randint(low, high, size, seed=None, *, dtype=None):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> print(ops.randint(1, 10, (2,3)))
         [[4 9 7]
          [9 1 2]]
@@ -1111,15 +1248,18 @@ def randint(low, high, size, seed=None, *, dtype=None):
     if dtype is None:
         dtype = mstype.int64
     elif dtype not in mstype.int_type:
-        raise ValueError(f"For 'randint', the 'dtype' must be an int type, but got {dtype}.")
+        raise ValueError(
+            f"For 'randint', the 'dtype' must be an int type, but got {dtype}.")
     if not isinstance(size, tuple):
-        raise ValueError(f"For 'randint', the input 'size' must be a tuple, but got {size}.")
+        raise ValueError(
+            f"For 'randint', the input 'size' must be a tuple, but got {size}.")
     if not isinstance(low, int) or isinstance(low, bool):
-        raise TypeError(f"For 'randint_like', 'low' must be an int, but got {type(low)}.")
+        raise TypeError(
+            f"For 'randint_like', 'low' must be an int, but got {type(low)}.")
     if not isinstance(high, int) or isinstance(high, bool):
-        raise TypeError(f"For 'randint_like', 'high' must be an int, but got {type(high)}.")
+        raise TypeError(
+            f"For 'randint_like', 'high' must be an int, but got {type(high)}.")
     seed1, seed2 = _get_seed(seed, 'randint')
-    cast_ = P.Cast()
     rand_op = P.UniformInt(seed1, seed2)
     rand_op = _set_prim_op_user_data(rand_op, "random_cache", False)
     low_ = Tensor(low, mstype.int32)
@@ -1164,20 +1304,23 @@ def randint_like(input, low, high, seed=None, *, dtype=None):
         [9 1 2]]
     """
     if not isinstance(input, Tensor):
-        raise TypeError(f"For 'randint_like', the 'input' must be a Tensor, but got {type(input)}")
+        raise TypeError(
+            f"For 'randint_like', the 'input' must be a Tensor, but got {type(input)}")
     if dtype is None:
         dtype = input.dtype
     if dtype not in mstype.int_type:
-        raise ValueError(f"For 'randint_like', the 'dtype' must be an int type, but got {dtype}.")
+        raise ValueError(
+            f"For 'randint_like', the 'dtype' must be an int type, but got {dtype}.")
     if not isinstance(low, int) or isinstance(low, bool):
-        raise TypeError(f"For 'randint_like', 'low' must be an int, but got {type(low)}.")
+        raise TypeError(
+            f"For 'randint_like', 'low' must be an int, but got {type(low)}.")
     if not isinstance(high, int) or isinstance(high, bool):
-        raise TypeError(f"For 'randint_like', 'high' must be an int, but got {type(high)}.")
+        raise TypeError(
+            f"For 'randint_like', 'high' must be an int, but got {type(high)}.")
     size = input.shape
     seed1, seed2 = _get_seed(seed, 'randint_like')
     rand_op = P.UniformInt(seed1, seed2)
     rand_op = _set_prim_op_user_data(rand_op, "random_cache", False)
-    cast_ = P.Cast()
     low_ = Tensor(low, mstype.int32)
     high_ = Tensor(high, mstype.int32)
     size_ = Tensor(size, mstype.int32)
@@ -1234,7 +1377,8 @@ def poisson(shape, mean, seed=None):
     """
     seed1, seed2 = _get_seed(seed, "poisson")
     random_poisson_op = P.Poisson(seed1, seed2)
-    random_poisson_op = _set_prim_op_user_data(random_poisson_op, "random_cache", False)
+    random_poisson_op = _set_prim_op_user_data(
+        random_poisson_op, "random_cache", False)
     value = random_poisson_op(shape, mean)
     return value
 
@@ -1329,50 +1473,51 @@ def multinomial(input, num_samples, replacement=True, seed=None):
         >>> # [[0 0 0 0 0 0 0 0 1 0]
         >>> #  [1 1 1 1 1 0 1 1 1 1]]
     """
-    shape = _get_cache_prim(P.Shape)()
-    reshape = _get_cache_prim(P.Reshape)()
-
     def _check_valid_dim(dim, name):
         if dim not in (1, 2):
-            raise ValueError(f"For '{name}', the dimension of inputs must be 1d or 2d, but got {dim}.")
+            raise ValueError(
+                f"For '{name}', the dimension of inputs must be 1d or 2d, but got {dim}.")
 
-    _check_valid_dim(len(shape(input)), "multinomial")
+    _check_valid_dim(len(shape_(input)), "multinomial")
     seed1, seed2 = _get_seed(seed, "multinomial")
     if not replacement:
-        if shape(input)[-1] < num_samples:
+        if shape_(input)[-1] < num_samples:
             const_utils.raise_value_error(f"For 'multinomial', the 'num_samples' must be less than "
                                           f"the last dimension of input without 'replacement', "
                                           f"but got 'num_samples': {num_samples} and "
                                           f"'replacement': {replacement}")
         n_dist = 1
-        if len(shape(input)) > 1:
-            n_dist = shape(input)[-2]
+        if len(shape_(input)) > 1:
+            n_dist = shape_(input)[-2]
         random_uniform_real = P.UniformReal(seed1, seed2)
-        random_cache_op = _set_prim_op_user_data(random_uniform_real, "random_cache", False)
-        random_uniform = random_cache_op((n_dist * shape(input)[-1],))
+        random_cache_op = _set_prim_op_user_data(
+            random_uniform_real, "random_cache", False)
+        random_uniform = random_cache_op((n_dist * shape_(input)[-1],))
         if n_dist != 1:
-            random_uniform = reshape(random_uniform, (n_dist, shape(input)[-1]))
-        real_div = _get_cache_prim(P.RealDiv)()
-        log = _get_cache_prim(P.Log)()
-        top_k = _get_cache_prim(P.TopK)()
+            random_uniform = reshape_(
+                random_uniform, (n_dist, shape_(input)[-1]))
 
-        vals = real_div(log(random_uniform), input + 1e-6)
-        _, indices = top_k(vals, num_samples)
+        vals = real_div_(log_(random_uniform), input + 1e-6)
+        _, indices = top_k_(vals, num_samples)
         return indices
     random_nomial = P.Multinomial(seed1, seed2)
-    random_nomial = _set_prim_op_user_data(random_nomial, "random_cache", False)
+    random_nomial = _set_prim_op_user_data(
+        random_nomial, "random_cache", False)
     return random_nomial(input, num_samples)
 
 
 def _check_shape(input_shape):
     """Check 'shape' value."""
     if not isinstance(input_shape, tuple):
-        const_utils.raise_type_error(f"Type of 'shape' must be tuple, but got: {type(input_shape)}")
+        const_utils.raise_type_error(
+            f"Type of 'shape' must be tuple, but got: {type(input_shape)}")
     for item in input_shape:
         if not isinstance(item, int):
-            const_utils.raise_type_error(f"Elements of 'shape' must be int, but got: {type(item)}")
+            const_utils.raise_type_error(
+                f"Elements of 'shape' must be int, but got: {type(item)}")
         if item < 1:
-            const_utils.raise_value_error(f"Elements of 'shape' must be positive int, but got: {item}")
+            const_utils.raise_value_error(
+                f"Elements of 'shape' must be positive int, but got: {item}")
     return True
 
 
@@ -1387,7 +1532,8 @@ def _check_param(op_name, param_name, param_value):
 __all__ = [
     'standard_laplace', 'random_categorical', 'uniform', 'standard_normal', 'random_gamma',
     'uniform_candidate_sampler', 'random_poisson', 'log_uniform_candidate_sampler', 'shuffle', 'choice_with_mask',
-    'normal', 'laplace', 'gamma', 'poisson', 'multinomial', 'rand', 'rand_like', 'randn', 'randn_like', 'randint',
-    'randint_like', 'multinomial_with_replacement', 'randperm'
+    'normal', 'laplace', 'gamma', 'poisson', 'multinomial', 'rand', 'rand_like',
+    'randn', 'randn_like',
+    'randint', 'randint_like', 'multinomial_with_replacement', 'randperm'
 ]
 __all__.sort()