mindspore 2.1.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.0__cp38-cp38-manylinux1_x86_64.whl

mindspore/ops/function/random_func.py

@@ -15,6 +15,7 @@
 """Defines parameter operators with functional form."""
 
 from __future__ import absolute_import
+import numpy as np
 
 from mindspore import context
 from mindspore.ops import operations as P
@@ -212,13 +213,13 @@ def multinomial_with_replacement(x, seed, offset, numsamples, replacement=False)
     """
     if not isinstance(seed, Tensor):
         if not isinstance(seed, int):
-            raise TypeError("For multinomial_with_replacement,",
-                            "the input[seed] must be int, but got {}.".format(type(seed)))
+            raise TypeError(f"For multinomial_with_replacement,",
+                            f"the input[seed] must be int, but got {type(seed)}.")
         seed = Tensor(seed, dtype=mstype.int64)
     if not isinstance(offset, Tensor):
         if not isinstance(offset, int):
-            raise TypeError("For multinomial_with_replacement,",
-                            "the input[offset] must be int, but got {}.".format(type(offset)))
+            raise TypeError(f"For multinomial_with_replacement,",
+                            f"the input[offset] must be int, but got {type(offset)}.")
         offset = Tensor(offset, dtype=mstype.int64)
     multinomial_with_replacement_ = P.MultinomialWithReplacement(numsamples=numsamples,
                                                                  replacement=replacement)
@@ -359,7 +360,7 @@ def uniform_candidate_sampler(true_classes,
     If unique=True, candidates are drawn without replacement, else unique=False with replacement.
 
     Args:
-        true_classes (Tensor): A Tensor. The target classes with a Tensor shape of :math:`(batch_size, num_true)` .
+        true_classes (Tensor): A Tensor. The target classes with a Tensor shape of :math:`(batch\_size, num\_true)` .
         num_true (int): The number of target classes in each training example.
         num_sampled (int): The number of classes to randomly sample. The sampled_candidates will have a shape
             of num_sampled. If unique=True, num_sampled must be less than or equal to range_max.
@@ -367,15 +368,17 @@ def uniform_candidate_sampler(true_classes,
         range_max (int): The number of possible classes, must be positive.
         seed (int): Used for random number generation, must be non-negative. If seed has a value of 0,
             the seed will be replaced with a randomly generated value. Default: ``0`` .
-        remove_accidental_hits (bool): Whether accidental hit is removed. Default: ``False`` .
+        remove_accidental_hits (bool): Whether accidental hit is removed.
+            Accidental hit is when one of the true classes matches one of the sample classes.
+            Set ``True`` to remove which accidentally sampling the true class as sample class. Default: ``False`` .
 
     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.
@@ -679,17 +682,15 @@ def normal(shape, mean, stddev, seed=None):
     Args:
         shape (tuple): The shape of random tensor to be generated.
           The format is :math:`(N,*)` where :math:`*` means, any number of additional dimensions.
-        mean (Union[Tensor, int, float]): The mean μ distribution parameter, which specifies the location of the peak,
-          with data type in [int8, int16, int32, int64, float16, float32].
-        stddev (Union[Tensor, int, float]): The deviation σ distribution parameter. It should be greater than 0,
-          with data type in [int8, int16, int32, int64, float16, float32].
+        mean (Union[Tensor, int, float]): The mean μ distribution parameter, which specifies the location of the peak.
+        stddev (Union[Tensor, int, float]): The deviation σ distribution parameter. It should be greater than 0.
         seed (int): Seed is used as entropy source for the Random number engines to generate pseudo-random numbers.
           The value 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
         of `mean` and `stddev`.
-        The dtype is float32.
+        The dtype is [float32, float64].
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -726,10 +727,6 @@ def normal(shape, mean, stddev, seed=None):
         mean = Tensor(mean)
     if not isinstance(stddev, Tensor):
         stddev = Tensor(stddev)
-    mean_dtype = F.dtype(mean)
-    stddev_dtype = F.dtype(stddev)
-    const_utils.check_type_valid(mean_dtype, mstype.int_type + (mstype.float16, mstype.float32), 'normal')
-    const_utils.check_type_valid(stddev_dtype, mstype.int_type + (mstype.float16, mstype.float32), 'normal')
     seed1, seed2 = _get_seed(seed, "normal")
     stdnormal = P.StandardNormal(seed1, seed2)
     stdnormal = _set_prim_op_user_data(stdnormal, "random_cache", False)
@@ -840,26 +837,24 @@ def gamma(shape, alpha, beta, seed=None):
         >>> alpha = Tensor(np.array([[3, 4], [5, 6]]), mindspore.float32)
         >>> beta = Tensor(np.array([1.0, 2]), mindspore.float32)
         >>> output = ops.gamma(shape, alpha, beta, seed=5)
-        >>> result = output.shape
         >>> print(output)
-        [[[ 2.2132034  5.8855834]]
-         [ 3.3981476  7.5805717]
-        [[ 3.3981476  7.5805717]]
-         [ 3.7190282 19.941492]
-        [[ 2.9512358  2.5969937]]
-         [ 3.786061   5.160872 ]]]
+        [[[ 2.2132034  5.8855834]
+          [ 3.8825176  8.6066265]]
+         [[ 3.3981476  7.5805717]
+          [ 3.7190282 19.941492 ]]
+         [[ 2.9512358  2.5969937]
+          [ 3.786061   5.160872 ]]]
         >>> # case 4: beta_shape is (2, 1), the output is different.
         >>> shape = (3, 1, 2)
         >>> alpha = Tensor(np.array([[3, 4], [5, 6]]), mindspore.float32)
         >>> beta = Tensor(np.array([[1.0], [2.0]]), mindspore.float32)
         >>> output = ops.gamma(shape, alpha, beta, seed=5)
-        >>> result = output.shape
         >>> print(output)
-        [[[ 5.6085486  7.8280783]]
-         [ 15.97684  16.116285]
-        [[ 1.8347423  1.713663]]
-         [ 3.2434065 15.667398]
-        [[ 4.2922077  7.3365674]]
+        [[[ 5.6085486  7.8280783]
+         [ 15.97684  16.116285]]
+        [[ 1.8347423  1.713663]
+         [ 3.2434065 15.667398]]
+        [[ 4.2922077  7.3365674]
          [ 5.3876944  13.159832 ]]]
     """
     seed1, seed2 = _get_seed(seed, "gamma")
@@ -1249,9 +1244,34 @@ def multinomial(input, num_samples, replacement=True, seed=None):
     Returns a tensor sampled from the multinomial probability distribution located in the corresponding
     row of the input tensor.
 
+    The polynomial distribution is a probability distribution that generalizes the binomial distribution formula to
+    multiple states. In the polynomial distribution, each event has a fixed probability, and the sum of these
+    probabilities is 1. The purpose of the `mindspore.ops.multinomial` interface is to perform `num_samples` sampling
+    on the input `input`, and the output tensor is the index of the input tensor for each sampling.
+    The values in `input` represent the probability of selecting the corresponding index for each sampling.
+
+    Here is an extreme example for better understanding. Suppose we have an input probability tensor with
+    values `Tensor([90 / 100, 10 / 100, 0], mindspore.float32)`, which means we can sample three indices,
+    namely index 0, index 1, and index 2, with probabilities of 90%, 10%, and 0%, respectively. We perform n samplings,
+    and the resulting sequence is the calculation result of the polynomial distribution, with a length equal to the
+    number of samplings.
+
+    In case 1 of the sample code, we perform two non-replacement samplings (`replacement` is `False`).
+    The calculation result is most likely `[0, 1]`, and less likely `[1, 0]`. Since the probability of selecting
+    index 0 is 90% for each sampling, the first result is most likely to be index 0. Since the probability of selecting
+    index 2 is 0, index 2 cannot appear in the sampling result. Therefore, the second result must be index 1,
+    and the resulting sequence is `[0, 1]`.
+
+    In case 2 of the sample code, we perform 10 replacement samplings (`replacement` is `True`).
+    As expected, about 90% of the sampling results are index 0.
+
+    In case 3 of the sample code, we extend the input to 2 dimensions, and the sampling results
+    in each dimension also match our sampling expectations.
+
     Note:
         The rows of input do not need to sum to one (in which case we use the values as weights),
-        but must be non-negative, finite and have a non-zero sum.
+        but must be non-negative, finite and have a non-zero sum. When using values as weights, it can be understood as
+        normalizing the input along the last dimension.
 
     Args:
         input (Tensor): The input tensor containing probabilities, must be 1 or 2 dimensions, with
@@ -1278,27 +1298,35 @@ def multinomial(input, num_samples, replacement=True, seed=None):
         >>> from mindspore import Tensor, ops
         >>> from mindspore import dtype as mstype
         >>> # case 1: The output is random, and the length of the output is the same as num_sample.
-        >>> input = Tensor([0, 9, 4, 0], mindspore.float32)
-        >>> output = ops.multinomial(input, 2)
-        >>> # print(output)
-        >>> # [1 2] or [2 1]
-        >>> # the case where the result is [2 1] in multiple times.
-        >>> # This is because the value corresponding to the index 1 is larger than the value of the index 2.
-        >>> print(len(output))
+        >>> # replacement is False.
+        >>> input1 = Tensor([90 / 100, 10 / 100, 0], mindspore.float32)
+        >>> input2 = Tensor([90, 10, 0], mindspore.float32)
+        >>> # input1 and input2 have the same meaning.
+        >>> output1 = ops.multinomial(input1, 2, replacement=False)
+        >>> output2 = ops.multinomial(input2, 2, replacement=False)
+        >>> # print(output1)
+        >>> # [0 1]
+        >>> # print(output2)
+        >>> # [0 1]
+        >>> print(len(output1))
+        2
+        >>> print(len(output2))
         2
         >>> # case 2: The output is random, and the length of the output is the same as num_sample.
-        >>> # replacement is False(Default).
-        >>> # If the extracted value is 0, the index value of 1 will be returned.
-        >>> input = Tensor([0, 9, 4, 0], mstype.float32)
-        >>> output = ops.multinomial(input, 4)
-        >>> print(output)
-        [1 1 2 1]
-        >>> # case 3: The output is random, num_sample == x_length = 4, and replacement is True,
-        >>> # Can extract the same elements。
-        >>> input = Tensor([0, 9, 4, 0], mstype.float32)
-        >>> output = ops.multinomial(input, 4, True)
-        >>> print(output)
-        [1 1 2 2]
+        >>> # replacement is True.
+        >>> output3 = ops.multinomial(input1, 10)
+        >>> # print(output3)
+        >>> # [0 0 1 0 0 0 0 0 0 0]
+        >>> print(len(output3))
+        10
+        >>> # case 3: The output is random, and the length of the output is the same as num_sample.
+        >>> # replacement is True.
+        >>> # rank is 2
+        >>> input4 = Tensor([[90, 10, 0], [10, 90, 0]], mstype.float32)
+        >>> output4 = ops.multinomial(input4, 10)
+        >>> # print(output4)
+        >>> # [[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)()
@@ -1318,7 +1346,9 @@ def multinomial(input, num_samples, replacement=True, seed=None):
         n_dist = 1
         if len(shape(input)) > 1:
             n_dist = shape(input)[-2]
-        random_uniform = _get_cache_prim(P.UniformReal)(seed1, seed2)((n_dist * shape(input)[-1],))
+        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],))
         if n_dist != 1:
             random_uniform = reshape(random_uniform, (n_dist, shape(input)[-1]))
         real_div = _get_cache_prim(P.RealDiv)()
@@ -1336,18 +1366,18 @@ def multinomial(input, num_samples, replacement=True, seed=None):
 def _check_shape(input_shape):
     """Check 'shape' value."""
     if not isinstance(input_shape, tuple):
-        const_utils.raise_type_error("Type of 'shape' must be tuple, but got: {}".format(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("Elements of 'shape' must be int, but got: {}".format(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("Elements of 'shape' must be positive int, but got: {}".format(item))
+            const_utils.raise_value_error(f"Elements of 'shape' must be positive int, but got: {item}")
     return True
 
 
 def _check_param(op_name, param_name, param_value):
     """Check type of param_value is Tensor, int, or float."""
-    if not isinstance(param_value, (Tensor, int, float)):
+    if not isinstance(param_value, (Tensor, int, float, np.ndarray)):
         const_utils.raise_type_error("For '{}', the type of '{}' must be Tensor, int, or float, "
                                      "but got: {}".format(op_name, param_name, type(param_value)))
     return True

mindspore/ops/function/sparse_func.py

@@ -423,7 +423,7 @@ def csr_to_dense(csr_tensor: CSRTensor) -> Tensor:
 
     valid_indices_dtype = [mstype.int32, mstype.int64]
     if row_pointers.dtype in valid_indices_dtype and col_indices.dtype in valid_indices_dtype:
-        if row_pointers.dtype == mstype.int64 or col_indices.dtype == mstype.int64:
+        if mstype.int64 in (row_pointers.dtype, col_indices.dtype):
             return csr_sparse_matrix_to_dense(dense_shape.astype(mstype.int64), batch_pointers.astype(mstype.int64),
                                               row_pointers.astype(mstype.int64), col_indices.astype(mstype.int64),
                                               values)

mindspore/ops/function/sparse_unary_func.py

@@ -383,7 +383,7 @@ def coo_relu(x: COOTensor) -> COOTensor:
     Args:
         x (COOTensor): Input COOTensor with shape :math:`(N, *)`, where :math:`*`
             means any number of additional dimensions. Its dtype is
-            `number <https://www.mindspore.cn/docs/en/r2.1/api_python/mindspore.html#mindspore.dtype>`_.
+            `number <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.html#mindspore.dtype>`_.
 
     Returns:
         COOTensor, has the same shape and dtype as the `x`.
@@ -1980,7 +1980,7 @@ def coo_acosh(x: COOTensor) -> COOTensor:
 
     .. math::
 
-        out_i = \cosh^{-1}(input_i)
+        y_i = \cosh^{-1}(x_i)
 
     .. warning::
         Given an input COOTensor x, the function computes inverse hyperbolic cosine of every element.
@@ -2414,14 +2414,16 @@ def coo_sin(x: COOTensor) -> COOTensor:
     return COOTensor(x.indices, math_func.sin(x.values), x.shape)
 
 
-__all__ = ["csr_cos", "csr_tan", "csr_exp", "csr_inv", "csr_relu", "csr_expm1", "csr_isfinite",
-           "csr_asin", "csr_sqrt", "csr_log", "csr_isnan", "csr_acos", "csr_floor", "csr_atan",
-           "csr_square", "csr_relu6", "csr_sinh", "csr_ceil", "csr_cosh", "csr_softsign",
-           "csr_log1p", "csr_round", "csr_tanh", "csr_asinh", "csr_neg", "csr_acosh", "csr_isinf",
-           "csr_atanh", "csr_sigmoid", "csr_abs", "csr_sin", "coo_cos", "coo_tan", "coo_exp",
-           "coo_inv", "coo_relu", "coo_expm1", "coo_isfinite", "coo_asin", "coo_sqrt", "coo_log",
-           "coo_isnan", "coo_acos", "coo_floor", "coo_atan", "coo_square", "coo_relu6", "coo_sinh",
-           "coo_ceil", "coo_cosh", "coo_softsign", "coo_log1p", "coo_round", "coo_tanh",
-           "coo_asinh", "coo_neg", "coo_acosh", "coo_isinf", "coo_atanh", "coo_sigmoid", "coo_abs",
-           "coo_sin"]
+__all__ = [
+    "csr_cos", "csr_tan", "csr_exp", "csr_inv", "csr_relu", "csr_expm1", "csr_isfinite",
+    "csr_asin", "csr_sqrt", "csr_log", "csr_isnan", "csr_acos", "csr_floor", "csr_atan",
+    "csr_square", "csr_relu6", "csr_sinh", "csr_ceil", "csr_cosh", "csr_softsign",
+    "csr_log1p", "csr_round", "csr_tanh", "csr_asinh", "csr_neg", "csr_acosh", "csr_isinf",
+    "csr_atanh", "csr_sigmoid", "csr_abs", "csr_sin", "coo_cos", "coo_tan", "coo_exp",
+    "coo_inv", "coo_relu", "coo_expm1", "coo_isfinite", "coo_asin", "coo_sqrt", "coo_log",
+    "coo_isnan", "coo_acos", "coo_floor", "coo_atan", "coo_square", "coo_relu6", "coo_sinh",
+    "coo_ceil", "coo_cosh", "coo_softsign", "coo_log1p", "coo_round", "coo_tanh",
+    "coo_asinh", "coo_neg", "coo_acosh", "coo_isinf", "coo_atanh", "coo_sigmoid", "coo_abs",
+    "coo_sin"
+]
 __all__.sort()

mindspore/ops/function/vmap_func.py

@@ -26,9 +26,11 @@ def vmap(fn, in_axes=0, out_axes=0):
 
     Vmap is pioneered by Jax and it removes the restriction of batch dimension on the operator, and provides a
     more convenient and unified operator expression. Moreover, it allows users to composite with other functional
-    modules such as :func:`mindspore.grad`, to improve the development efficiency. In addition, the vectorizing
-    map does not execute loops outside the function, but sinks loops into the primitive operations of the function
-    for better performance. When combined with `Graph Kernel Fusion`, operational efficiency would be further improved.
+    modules such as :func:`mindspore.grad`, to improve the development efficiency, please refer to the
+    `Automatic Vectorization (Vmap) <https://www.mindspore.cn/tutorials/experts/en/r2.2/vmap/vmap.html>`_ tutorial
+    for more detail. In addition, the vectorizing map does not execute loops outside the function, but sinks loops
+    into the primitive operations of the function for better performance. When combined with `Graph Kernel Fusion`,
+    operational efficiency would be further improved.
 
     .. warning::
         This is an experimental API that is subject to change or deletion.
@@ -36,8 +38,7 @@ def vmap(fn, in_axes=0, out_axes=0):
     Note:
         1. The power of vmap comes from the implementation of VmapRules of primitives. Although we have designed a
         generalized rule for user custom operators, we can not guarantee that it works well for all operators,
-        please be aware the risk of use. If you want to achieve a better performance, please refer to the tutorial to
-        implement the specific VmapRule for the custom operator, which won't take too much time.
+        unknown exceptions may occur, please be aware the risk of use.
         2. When calling the random number generation methods within the scope of vmap, the same random number is
         generated among vector functions each time. If you expect each vector branch to use different random numbers,
         you need to generate batch random numbers externally in advance and then transfer them to vmap.

mindspore/ops/functional.py

@@ -20,6 +20,7 @@ from mindspore.common._register_for_tensor import tensor_operator_registry
 from mindspore.ops import _constants
 from mindspore.ops.function import *
 from mindspore.ops.function.array_func import narrow, flatten
+from mindspore.ops.function.math_func import all
 from mindspore.ops import operations as P
 from mindspore.ops.operations import array_ops
 from mindspore.ops.primitive import Primitive
@@ -122,15 +123,16 @@ reduced_shape = Primitive("reduced_shape")
 # shape_mul:input must be shape multiply elements in tuple(shape)
 shape_mul = _sequence_ops.shape_mul()
 
+tensor_operator_registry.register('tuple_to_tensor', _sequence_ops.TupleToTensor)
 tensor_operator_registry.register('add', P.Add)
 tensor_operator_registry.register('addr', addr)
 tensor_operator_registry.register('addcdiv', P.Addcdiv)
 tensor_operator_registry.register('addcmul', P.Addcmul)
-tensor_operator_registry.register('all', P.ReduceAll)
+tensor_operator_registry.register('all', all)
 tensor_operator_registry.register('angle', angle)
 tensor_operator_registry.register('any', P.ReduceAny)
 tensor_operator_registry.register('atan2', atan2)
-tensor_operator_registry.register('abs', P.Abs)
+tensor_operator_registry.register('abs', abs)
 tensor_operator_registry.register('baddbmm', baddbmm)
 tensor_operator_registry.register('geqrf', geqrf)
 tensor_operator_registry.register('histc', histc)
@@ -142,6 +144,7 @@ tensor_operator_registry.register('slogdet', slogdet)
 tensor_operator_registry.register('trace', trace)
 tensor_operator_registry.register('tril', tril)
 tensor_operator_registry.register('chunk', chunk)
+tensor_operator_registry.register('count_nonzero', count_nonzero)
 tensor_operator_registry.register('sqrt', sqrt)
 tensor_operator_registry.register('square', square)
 tensor_operator_registry.register('sub', sub)
@@ -163,10 +166,10 @@ tensor_operator_registry.register('negative', neg)
 tensor_operator_registry.register('amin', amin)
 tensor_operator_registry.register('amax', amax)
 tensor_operator_registry.register('aminmax', aminmax)
-tensor_operator_registry.register('mean', P.ReduceMean)
+tensor_operator_registry.register('mean', mean)
 tensor_operator_registry.register('prod', prod)
 tensor_operator_registry.register('round', P.Round)
-tensor_operator_registry.register('reshape', P.Reshape)
+tensor_operator_registry.register('reshape', reshape)
 tensor_operator_registry.register('reverse', P.ReverseV2)
 tensor_operator_registry.register('reverse_sequence', P.ReverseSequence)
 tensor_operator_registry.register('xlogy', P.Xlogy)
@@ -176,8 +179,8 @@ tensor_operator_registry.register('broadcast_to', P.BroadcastTo)
 tensor_operator_registry.register('matmul', matmul)
 tensor_operator_registry.register('inner', inner)
 tensor_operator_registry.register('xdivy', P.Xdivy)
-tensor_operator_registry.register('argmax', P.Argmax)
-tensor_operator_registry.register('argmin', P.Argmin)
+tensor_operator_registry.register('argmax', argmax)
+tensor_operator_registry.register('argmin', argmin)
 tensor_operator_registry.register('cumsum', P.CumSum)
 tensor_operator_registry.register('cummin', cummin)
 tensor_operator_registry.register('cummax', cummax)
@@ -216,10 +219,11 @@ tensor_operator_registry.register('logdet', logdet)
 tensor_operator_registry.register('log_matrix_determinant', log_matrix_determinant)
 tensor_operator_registry.register('matrix_determinant', matrix_determinant)
 tensor_operator_registry.register('ceil', P.Ceil)
-tensor_operator_registry.register('fill', P.Fill)
-tensor_operator_registry.register('tile', P.Tile)
+tensor_operator_registry.register('fillv2', P.FillV2)
+tensor_operator_registry.register('tile', tile)
 tensor_operator_registry.register('logit', logit)
-tensor_operator_registry.register('sum', P.ReduceSum)
+tensor_operator_registry.register('sum', sum)
+tensor_operator_registry.register('reducesum', P.ReduceSum)
 tensor_operator_registry.register('split', split)
 tensor_operator_registry.register('tensor_split', tensor_split)
 tensor_operator_registry.register('vsplit', vsplit)
@@ -264,7 +268,7 @@ tensor_operator_registry.register('standard_normal', P.StandardNormal)
 tensor_operator_registry.register('sigmoid', P.Sigmoid)
 tensor_operator_registry.register('median', Median)
 tensor_operator_registry.register('tanh', tanh)
-tensor_operator_registry.register('exp', P.Exp)
+tensor_operator_registry.register('exp', exp)
 tensor_operator_registry.register('addbmm', addbmm)
 tensor_operator_registry.register('addmm', addmm)
 tensor_operator_registry.register('addmv', addmv)
@@ -284,6 +288,7 @@ tensor_operator_registry.register('ldexp', ldexp)
 tensor_operator_registry.register('clamp', clamp)
 tensor_operator_registry.register('fold', fold)
 tensor_operator_registry.register('unfold', unfold)
+tensor_operator_registry.register('diagonal', diagonal)
 tensor_operator_registry.register('diagonal_scatter', diagonal_scatter)
 tensor_operator_registry.register('index_add', index_add)
 tensor_operator_registry.register('greater', greater)