• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
2#
3# Copyright 2020-2021 Huawei Technologies Co., Ltd
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16# ============================================================================
17"""standard_method"""
18
19from dataclasses import dataclass
20
21from mindspore import Tensor, Parameter
22from mindspore import dtype as mstype
23
24from ..._checkparam import Validator as validator
25from ...ops import functional as F
26from ...ops import operations as P
27from ...ops.composite import tail, core, MultitypeFuncGraph, env_get, hyper_add, \
28    zeros_like, ones_like, repeat_elements
29from ...ops.composite.base import _append
30from ...ops.composite.multitype_ops import _constexpr_utils as const_utils
31from ...ops.composite.multitype_ops import _compile_utils as compile_utils
32from ...ops.primitive import constexpr
33
34
35__all__ = ['MultitypeFuncGraph', 'env_get', 'hyper_add', 'zeros_like', 'ones_like']
36
37shape_ = P.Shape()
38dtype_ = P.DType()
39abs_ = P.Abs()
40ndim_ = P.Rank()
41cumsum_ = P.CumSum()
42size_op_ = P.Size()
43_reduce_sum_default = P.ReduceSum()
44_reduce_sum_keepdims = P.ReduceSum(True)
45_mean_keepdims = P.ReduceMean(True)
46
47itemsize_map = {mstype.bool_: 1, mstype.int8: 1, mstype.uint8: 1,
48                mstype.float16: 2, mstype.int16: 2, mstype.uint16: 2,
49                mstype.float32: 4, mstype.int32: 4, mstype.uint32: 4,
50                mstype.float64: 8, mstype.int64: 8, mstype.uint64: 8}
51
52nan_tensor = Tensor(float('nan'), dtype=mstype.float32)
53
54
55def mean(x, axis=(), keep_dims=False):
56    """
57    Reduces a dimension of a tensor by averaging all elements in the dimension.
58
59    Args:
60        axis (Union[None, int, tuple(int), list(int)]): Dimensions of reduction,
61            when axis is None or empty tuple, reduce all dimensions. Default: ().
62        keep_dims (bool): Whether to keep the reduced dimensions. Default: False.
63
64    Returns:
65        Tensor, has the same data type as input tensor.
66
67    Supported Platforms:
68        ``Ascend`` ``GPU`` ``CPU``
69
70    Examples:
71        >>> import numpy as np
72        >>> from mindspore import Tensor
73        >>> input_x = Tensor(np.array([1, 2, 3], dtype=np.float32))
74        >>> output = input_x.mean()
75        >>> print(output)
76        2.0
77    """
78    if axis is None:
79        axis = ()
80    reduce_mean = P.ReduceMean(keep_dims)
81    return reduce_mean(x, axis)
82
83
84def all_(x, axis=(), keep_dims=False):
85    """
86    Check all array elements along a given axis evaluate to True.
87
88    Args:
89        x (Tensor): A Tensor to be reduced.
90        axis (Union[None, int, tuple(int)): Dimensions of reduction.
91        keep_dims (bool): Whether to keep the reduced dimensions.
92
93    Returns:
94        Tensor, has the same data type as x.
95    """
96
97    if axis is None:
98        axis = ()
99    reduce_all = P.ReduceAll(keep_dims)
100    return reduce_all(x, axis)
101
102
103def any_(x, axis=(), keep_dims=False):
104    """
105    Check any array element along a given axis evaluate to True.
106
107    Args:
108        x (Tensor): A Tensor to be reduced.
109        axis (Union[None, int, tuple(int)): Dimensions of reduction.
110        keep_dims (bool): Whether to keep the reduced dimensions.
111
112    Returns:
113        Tensor, has the same data type as x.
114    """
115    if axis is None:
116        axis = ()
117    reduce_any = P.ReduceAny(keep_dims)
118    return reduce_any(x, axis)
119
120
121def size_(x):
122    """
123    Return the number of elements in tensor `x`.
124
125    Note:
126        To strictly follow Numpy's behaviour, return 1 for tensor scalar.
127
128    Args:
129        x (Tensor): Input tensor.
130
131    Returns:
132        size(int).
133    """
134    if not shape_(x):
135        return size_op_(x) + 1
136    return size_op_(x)
137
138
139def itemsize_(x):
140    """
141    Return length of one tensor element in bytes.
142
143    Args:
144        x (Tensor): Input tensor.
145
146    Returns:
147        itemsize(int).
148    """
149    return get_itemsize(x.dtype)
150
151
152def nbytes_(x):
153    """
154    Return total number of bytes taken by the tensor.
155
156    Args:
157        x (Tensor): Input tensor.
158
159    Returns:
160        nbytes(int).
161    """
162    return itemsize_(x) * F.shape_mul(shape_(x))
163
164
165def strides_(x):
166    """
167    Return the tuple of bytes to step in each dimension when traversing a tensor.
168
169    Args:
170        x (Tensor): Input tensor.
171
172    Returns:
173        strides (tuple[int]).
174    """
175    strides = ()
176    ndim = P.Rank()(x)
177    tensor_shape = shape_(x)
178    for i in F.make_range(0, ndim):
179        stride = itemsize_(x)
180        for j in F.make_range(i + 1, ndim):
181            stride *= tensor_shape[j]
182        strides += (stride,)
183    return strides
184
185
186def astype(x, dtype, copy=True): # pylint: disable=redefined-outer-name
187    """
188    Return a copy of the tensor, casted to a specified type.
189
190    Args:
191        dtype (Union[:class:`mindspore.dtype`, str]): Designated tensor dtype, can be in format
192            of :class:`mindspore.dtype.float32` or `float32`.
193            Default: :class:`mindspore.dtype.float32`.
194        copy (bool, optional): By default, astype always returns a newly allocated
195            tensor. If this is set to false, the input tensor is returned instead
196            of a copy if possible. Default: True.
197
198    Returns:
199        Tensor, with the designated dtype.
200
201    Raises:
202        TypeError: If `dtype` has types not specified above, or values cannot be understood.
203
204    Supported Platforms:
205        ``Ascend`` ``GPU`` ``CPU``
206
207    Examples:
208        >>> import numpy as np
209        >>> from mindspore import Tensor
210        >>> x = Tensor(np.ones((1,2,2,1), dtype=np.float32))
211        >>> x = x.astype("int32")
212        >>> print(x.dtype)
213        Int32
214    """
215    dtype = check_astype_dtype_const(dtype)
216    if not copy and dtype == x.dtype:
217        return x
218    return F.cast(x, dtype)
219
220
221def transpose(x, *axis):
222    r"""
223    Return a view of the tensor with axes transposed.
224
225    For a 1-D tensor this has no effect, as a transposed vector is simply the
226    same vector. For a 2-D tensor, this is a standard matrix transpose. For a
227    n-D tensor, if axes are given, their order indicates how the axes are permuted.
228    If axes are not provided and tensor.shape = (i[0], i[1],...i[n-2], i[n-1]),
229    then tensor.transpose().shape = (i[n-1], i[n-2], ... i[1], i[0]).
230
231    Args:
232        axes(Union[None, tuple(int), list(int), int], optional): If axes is None or
233            blank, tensor.transpose() will reverse the order of the axes. If axes is tuple(int)
234            or list(int), tensor.transpose() will transpose the tensor to the new axes order.
235            If axes is int, this form is simply intended as a convenience alternative to the
236            tuple/list form.
237
238    Returns:
239        Tensor, has the same dimension as input tensor, with axes suitably permuted.
240
241    Raises:
242        TypeError: If input arguments have types not specified above.
243        ValueError: If the number of `axes` is not euqal to a.ndim.
244
245    Supported Platforms:
246        ``Ascend`` ``GPU`` ``CPU``
247
248    Examples:
249        >>> import numpy as np
250        >>> from mindspore import Tensor
251        >>> x = Tensor(np.ones((1,2,3), dtype=np.float32))
252        >>> x = x.transpose()
253        >>> print(x.shape)
254        (3, 2, 1)
255    """
256    ndim = F.rank(x)
257    perm = check_transpose_axis_const(axis, ndim)
258    return F.transpose(x, perm)
259
260
261# `tensor.T` is used as a property in graph mode
262T_ = transpose
263
264
265def reshape(x, *shape):
266    """
267    Give a new shape to a tensor without changing its data.
268
269    Args:
270        shape(Union[int, tuple(int), list(int)]): The new shape should be compatible
271            with the original shape. If an integer, then the result will be a 1-D
272            array of that length. One shape dimension can be -1. In this case, the
273            value is inferred from the length of the array and remaining dimensions.
274
275    Returns:
276        Tensor, with new specified shape.
277
278    Raises:
279        TypeError: If new_shape is not integer, list or tuple, or `x` is not tensor.
280        ValueError: If new_shape is not compatible with the original shape.
281
282    Supported Platforms:
283        ``Ascend`` ``GPU`` ``CPU``
284
285    Examples:
286        >>> from mindspore import Tensor
287        >>> from mindspore import dtype as mstype
288        >>> x = Tensor([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]], dtype=mstype.float32)
289        >>> output = x.reshape((3, 2))
290        >>> print(output)
291        [[-0.1  0.3]
292        [ 3.6  0.4]
293        [ 0.5 -3.2]]
294    """
295    new_shape = check_reshape_shp_const(shape)
296    return F.reshape(x, new_shape)
297
298
299def ravel(x):
300    """
301    Return a contiguous flattened tensor.
302
303    Returns:
304        Tensor, a 1-D tensor, containing the same elements of the input.
305
306    Supported Platforms:
307        ``Ascend`` ``GPU`` ``CPU``
308
309    Examples:
310        >>> import numpy as np
311        >>> from mindspore import Tensor
312        >>> x = Tensor(np.ones((2,3,4), dtype=np.float32))
313        >>> output = x.ravel()
314        >>> print(output.shape)
315        (24,)
316    """
317    return reshape(x, (-1,))
318
319
320def flatten(x, order='C'):
321    r"""
322    Return a copy of the tensor collapsed into one dimension.
323
324    Args:
325        order (str, optional): Can choose between 'C' and 'F'. 'C' means to
326            flatten in row-major (C-style) order. 'F' means to flatten in column-major
327            (Fortran-style) order. Only 'C' and 'F' are supported. Default: 'C'.
328
329    Returns:
330        Tensor, has the same data type as input.
331
332    Supported Platforms:
333        ``Ascend`` ``GPU`` ``CPU``
334
335    Raises:
336        TypeError: If `order` is not string type.
337        ValueError: If `order` is string type, but not 'C' or 'F'.
338
339    Examples:
340        >>> import numpy as np
341        >>> from mindspore import Tensor
342        >>> x = Tensor(np.ones((2,3,4), dtype=np.float32))
343        >>> output = x.flatten()
344        >>> print(output.shape)
345        (24,)
346    """
347    order = check_flatten_order_const(order)
348    if order == 'C':
349        return F.reshape(x, (-1,))
350
351    perm = F.make_range(0, F.rank(x))
352    new_order = F.tuple_reversed(perm)
353    return F.reshape(F.transpose(x, new_order), (-1,))
354
355
356def swapaxes(x, axis1, axis2):
357    """
358    Interchange two axes of a tensor.
359
360    Args:
361        axis1 (int): First axis.
362        axis2 (int): Second axis.
363
364    Returns:
365        Transposed tensor, has the same data type as the input.
366
367    Raises:
368        TypeError: If `axis1` or `axis2` is not integer.
369        ValueError: If `axis1` or `axis2` is not in the range of :math:`[-ndim, ndim-1]`.
370
371    Supported Platforms:
372        ``Ascend`` ``GPU`` ``CPU``
373
374    Examples:
375        >>> import numpy as np
376        >>> from mindspore import Tensor
377        >>> x = Tensor(np.ones((2,3,4), dtype=np.float32))
378        >>> output = x.swapaxes(0, 2)
379        >>> print(output.shape)
380        (4,3,2)
381    """
382    axis1, axis2 = check_swapaxes_axis_const((axis1, axis2), x.ndim)
383
384    if axis1 == axis2:
385        return x
386    if axis1 > axis2:
387        axis1, axis2 = axis2, axis1
388
389    perm = F.make_range(0, x.ndim)
390    new_perm = None
391    if axis2 + 1 < x.ndim:
392        new_perm = perm[0:axis1] + perm[axis2:axis2 + 1] + \
393            perm[axis1 + 1:axis2] + perm[axis1:axis1 + 1] + perm[axis2 + 1:]
394    else:
395        new_perm = perm[0:axis1] + perm[axis2:axis2 + 1] + \
396            perm[axis1 + 1:axis2] + perm[axis1:axis1 + 1]
397
398    return F.transpose(x, new_perm)
399
400
401def squeeze(x, axis=None):
402    """
403    Remove single-dimensional entries from the shape of a tensor.
404
405    Args:
406        axis (Union[None, int, list(int), tuple(int)], optional): Default is None.
407
408    Returns:
409        Tensor, with all or a subset of the dimensions of length 1 removed.
410
411    Raises:
412        TypeError: If input arguments have types not specified above.
413        ValueError: If specified axis has shape entry :math:`> 1`.
414
415    Supported Platforms:
416        ``Ascend`` ``GPU`` ``CPU``
417
418    Examples:
419        >>> import numpy as np
420        >>> from mindspore import Tensor
421        >>> x = Tensor(np.ones((1,2,2,1), dtype=np.float32))
422        >>> x = x.squeeze()
423        >>> print(x.shape)
424        (2, 2)
425    """
426    shape = F.shape(x)
427    if axis is None:
428        return F.squeeze(x)
429    # yield squeezed shape based on the axes
430    new_shape = prepare_shape_for_squeeze_const(shape, axis)
431    return F.reshape(x, new_shape)
432
433
434def argmax(x, axis=None):
435    """
436    Returns the indices of the maximum values along an axis.
437
438    Args:
439        axis (int, optional): By default, the index is into
440            the flattened array, otherwise along the specified axis.
441
442    Returns:
443        Tensor, array of indices into the array. It has the same
444        shape as a.shape with the dimension along axis removed.
445
446    Raises:
447        ValueError: if axis is out of range.
448
449    Supported Platforms:
450        ``Ascend`` ``GPU`` ``CPU``
451
452    Examples:
453        >>> import numpy as np
454        >>> from mindspore import Tensor
455        >>> a = Tensor(np.arange(10, 16).reshape(2, 3).astype("float32"))
456        >>> print(a.argmax())
457        5
458    """
459    # P.Argmax only supports float
460    x = x.astype(mstype.float32)
461    if axis is None:
462        x = ravel(x)
463        axis = 0
464    else:
465        axis = check_axis_in_range_const(axis, F.rank(x))
466    return P.Argmax(axis)(x)
467
468
469def argmin(x, axis=None):
470    """
471    Returns the indices of the minimum values along an axis.
472
473    Args:
474        a (Union[int, float, bool, list, tuple, Tensor]): Input array.
475        axis (int, optional): By default, the index is into
476            the flattened array, otherwise along the specified axis.
477
478    Returns:
479        Tensor, array of indices into the array. It has the same
480        shape as a.shape with the dimension along axis removed.
481
482    Raises:
483        ValueError: if axis is out of range.
484
485    Supported Platforms:
486        ``Ascend`` ``GPU`` ``CPU``
487
488    Examples:
489        >>> import numpy as np
490        >>> from mindspore import Tensor
491        >>> a = Tensor(np.arange(10, 16).reshape(2, 3).astype("float32"))
492        >>> print(a.argmin())
493        0
494    """
495    # P.Argmax only supports float
496    x = x.astype(mstype.float32)
497    if axis is None:
498        x = ravel(x)
499        axis = 0
500    else:
501        axis = check_axis_in_range_const(axis, F.rank(x))
502    # P.Argmin is currently not supported
503    return P.Argmax(axis)(F.neg_tensor(x))
504
505
506def cumsum(x, axis=None, dtype=None):
507    """
508    Returns the cumulative sum of the elements along a given axis.
509
510    Note:
511        If ``x.dtype`` is :class:`int8`, :class:`int16` or :class:`bool`, the result
512        `dtype` will be elevated to :class:`int32`, :class:`int64` is not supported.
513
514    Args:
515        x (Tensor): Input tensor.
516        axis (int, optional): Axis along which the cumulative sum is computed. The
517            default (None) is to compute the cumsum over the flattened array.
518        dtype (:class:`mindspore.dtype`, optional): If not specified, stay the same as original,
519            tensor, unless it has an integer dtype with a precision less than :class:`float32`.
520            In that case, :class:`float32` is used.
521
522    Returns:
523        Tensor.
524
525    Supported Platforms:
526        ``Ascend`` ``GPU`` ``CPU``
527
528    Examples:
529        >>> import numpy as np
530        >>> from mindspore import Tensor
531        >>> a = Tensor(np.ones((3,3)).astype("float32"))
532        >>> output = a.cumsum(axis=0)
533        >>> print(output)
534        [[1. 1. 1.]
535        [2. 2. 2.]
536        [3. 3. 3.]]
537    """
538    original_dtype = x.dtype
539    # If original tensor is int, and has precision less then int32, convert
540    # to int32
541    if x.dtype in (mstype.bool_, mstype.int8, mstype.int16, mstype.uint8, mstype.int16):
542        x = x.astype(mstype.int32)
543    if axis is None:
544        x = x.ravel()
545        axis = 0
546    check_axis_in_range_const(axis, x.ndim)
547    if dtype is not None and original_dtype != dtype:
548        return cumsum_(x, axis).astype(dtype, copy=False)
549    return cumsum_(x, axis)
550
551
552def copy(x):
553    """
554    Returns a copy of the tensor.
555
556    Note:
557        The current implementation does not support `order` argument.
558
559    Args:
560        x (Tensor): Input tensor.
561
562    Returns:
563        Copied tensor.
564
565    Supported Platforms:
566        ``Ascend`` ``GPU`` ``CPU``
567
568    Examples:
569        >>> import numpy as np
570        >>> from mindspore import Tensor
571        >>> a = Tensor(np.ones((3,3)).astype("float32"))
572        >>> output = a.copy()
573        >>> print(output)
574        [[1. 1. 1.]
575        [1. 1. 1.]
576        [1. 1. 1.]]
577    """
578    if x.size == 0:
579        return x
580    origin_dtype = x.dtype
581    if origin_dtype == mstype.bool_:
582        return F.logical_not(F.logical_not(x))
583    if origin_dtype != mstype.float64:
584        x = x.astype(mstype.float32)
585    x = x / 1.0
586    x = x.astype(origin_dtype)
587    return x
588
589
590def max(x, axis=None, keepdims=False, initial=None, where=True): # pylint: disable=redefined-builtin
591    """
592    Returns the maximum of a tensor or maximum along an axis.
593
594    Args:
595        x (Tensor): Input Tensor.
596        axis (None or int or tuple of ints, optional): defaults to None. Axis or
597            axes along which to operate. By default, flattened input is used. If
598            this is a tuple of ints, the maximum is selected over multiple axes,
599            instead of a single axis or all the axes as before.
600        keepdims (boolean, optional): defaults to False.
601            If this is set to True, the axes which are reduced are left in the
602            result as dimensions with size one. With this option, the result will
603            broadcast correctly against the input array.
604        initial (scalar, optional):
605            The minimum value of an output element. Must be present to allow
606            computation on empty slice.
607        where (boolean Tensor, optional): defaults to True.
608            A boolean array which is broadcasted to match the dimensions of array,
609            and selects elements to include in the reduction. If non-default value
610            is passed, initial must also be provided.
611
612    Returns:
613        Tensor or scalar, maximum of input tensor. If `axis` is None, the result is a scalar
614        value. If `axis` is given, the result is an array of dimension ``a.ndim - 1``.
615
616    Raises:
617        TypeError: if the input is not a tensor.
618
619    Supported Platforms:
620            ``Ascend`` ``GPU`` ``CPU``
621
622    Examples:
623        >>> import numpy as np
624        >>> from mindspore import Tensor
625        >>> import mindspore.numpy as np
626        >>> a = Tensor(np.arange(4).reshape((2,2)).astype('float32'))
627        >>> output = a.max()
628        >>> print(output)
629        3.0
630    """
631    return compile_utils.reduce_(x, P.ReduceMax(keepdims), cmp_fn=F.maximum,
632                                 axis=axis, keepdims=keepdims, initial=initial, where=where)
633
634
635def min(x, axis=None, keepdims=False, initial=None, where=True): # pylint: disable=redefined-builtin
636    """
637    Returns the minimum of a tensor or minimum along an axis.
638
639    Args:
640        a (Tensor): Input data.
641        axis (None or int or tuple of ints, optional): defaults to None. Axis or
642            axes along which to operate. By default, flattened input is used. If
643            this is a tuple of ints, the minimum is selected over multiple axes,
644            instead of a single axis or all the axes as before.
645        keepdims (boolean, optional): defaults to False.
646            If this is set to True, the axes which are reduced are left in the
647            result as dimensions with size one. With this option, the result will
648            broadcast correctly against the input array.
649        initial (scalar, optional):
650            The maximum value of an output element. Must be present to allow
651            computation on empty slice.
652        where (boolean Tensor, optional): defaults to True.
653            A boolean array which is broadcasted to match the dimensions of array,
654            and selects elements to include in the reduction. If non-default value
655            is passed, initial must also be provided.
656
657    Returns:
658        Tensor or scalar, minimum of `a`. If axis is None, the result is a scalar
659        value. If `axis` is given, the result is an array of dimension ``a.ndim - 1``.
660
661    Raises:
662        TypeError: if the input is not a tensor.
663
664    Supported Platforms:
665        ``Ascend`` ``GPU`` ``CPU``
666
667    Examples:
668        >>> import numpy as np
669        >>> from mindspore import Tensor
670        >>> import mindspore.numpy as np
671        >>> a = Tensor(np.arange(4).reshape((2,2)).astype('float32'))
672        >>> output = a.min()
673        >>> print(output)
674        0.0
675    """
676    return compile_utils.reduce_(x, P.ReduceMin(keepdims), cmp_fn=F.minimum,
677                                 axis=axis, keepdims=keepdims, initial=initial, where=where)
678
679
680def resize(x, *new_shape):
681    """
682    Changes shape and size of array in-place.
683
684    Note:
685        Instead of changing the size of the input array and returns nothing as in numpy,
686        this method returns a new Tensor with the input size.
687        Numpy argument `refcheck` is not supported.
688
689    Args:
690        new_shape (Union[ints, tuple of ints]): Shape of resized array.
691
692    Returns:
693        Tensor.
694
695    Supported Platforms:
696        ``Ascend`` ``GPU`` ``CPU``
697
698    Examples:
699        >>> from mindspore import numpy as np
700        >>> x = np.array([[0, 1], [2, 3]])
701        >>> x = x.resize(2, 3)
702        >>> print(x)
703        [[0 1 2]
704        [3 0 0]]
705    """
706    if not new_shape:
707        return x
708    if len(new_shape) == 1:
709        if isinstance(new_shape[0], tuple):
710            new_shape = new_shape[0]
711    flattened = x.ravel()
712    cur_size = F.shape_mul(x.shape)
713    new_size = F.shape_mul(new_shape)
714    diff_size = new_size - cur_size
715    if diff_size > 0:
716        pad_val = F.fill(x.dtype, (diff_size,), 0)
717        res = P.Concat()((flattened, pad_val))
718    else:
719        res = flattened[:new_size]
720    return res.reshape(new_shape)
721
722
723def diagonal(x, offset=0, axis1=0, axis2=1):
724    """
725    Returns specified diagonals.
726
727    Args:
728        offset (int, optional): Offset of the diagonal from the main diagonal.
729            Can be positive or negative. Defaults to main diagonal.
730        axis1 (int, optional): Axis to be used as the first axis of the 2-D
731            sub-arrays from which the diagonals should be taken. Defaults to
732            first axis (0).
733        axis2 (int, optional): Axis to be used as the second axis of the 2-D
734            sub-arrays from which the diagonals should be taken. Defaults to
735            second axis.
736
737    Returns:
738        Tensor, if `a` is 2-D, then `a` 1-D array containing the diagonal.
739
740    Raises:
741        ValueError: if the input tensor has less than two dimensions.
742
743    Supported Platforms:
744        ``Ascend`` ``GPU`` ``CPU``
745
746    Examples:
747        >>> import mindspore.numpy as np
748        >>> a = np.arange(4).reshape(2,2)
749        >>> print(a)
750        [[0 1]
751        [2 3]]
752        >>> output = a.diagonal()
753        >>> print(output)
754        [0 3]
755    """
756    ndim = x.ndim
757    if ndim < 2:
758        const_utils.raise_value_error('diagonal requires an array of at least two dimensions')
759    dtype = x.dtype
760
761    axes = check_axis_valid((axis1, axis2), ndim)
762    perm = ()
763    for i in range(ndim):
764        if i not in axes:
765            perm += (i,)
766    perm += axes
767    x = x.transpose(perm)
768
769    shape = x.shape
770    n, m = shape[-2:]
771
772    e = F.eye(n, m, dtype)
773    if offset >= m or offset <= -n:
774        e = F.fill(dtype, (n, m), 0)
775    elif offset != 0:
776        e = e.astype(mstype.float32)
777        if offset > 0:
778            e_left = F.fill(dtype, (n, offset), 0)
779            e_right = e[..., 0:m-offset:1]
780            e = P.Concat(1)((e_left, e_right)).astype(dtype)
781        elif offset < 0:
782            e_upper = F.fill(dtype, (-offset, m), 0)
783            e_lower = e[0:n+offset:1, ...]
784            e = P.Concat(0)((e_upper, e_lower)).astype(dtype)
785    e = P.BroadcastTo(shape)(e)
786
787    prod = F.tensor_mul(x, e)
788    res = F.reduce_sum(prod.astype(mstype.float32), -1)
789
790    begin = ()
791    for i in range(ndim-2):
792        begin += (0,)
793    last_dim_begin = max_(0, -offset)
794    begin += (last_dim_begin,)
795    size = res.shape[:-1]
796    last_dim_end = min_(
797        shape[-2], max_(0, shape[-1] - offset)) - last_dim_begin
798    if last_dim_end <= 0:
799        return empty_compile(dtype, (0,))
800    size += (last_dim_end,)
801    res = F.tensor_slice(res, begin, size)
802    return res.astype(dtype)
803
804
805def trace(x, offset=0, axis1=0, axis2=1, dtype=None):
806    """
807    Returns the sum along diagonals of the array.
808
809    Args:
810        offset (int, optional): Offset of the diagonal from the main diagonal.
811            Can be positive or negative. Defaults to main diagonal.
812        axis1 (int, optional): Axis to be used as the first axis of the 2-D
813            sub-arrays from which the diagonals should be taken. Defaults to
814            first axis (0).
815        axis2 (int, optional): Axis to be used as the second axis of the 2-D
816            sub-arrays from which the diagonals should be taken. Defaults to
817            second axis.
818        dtype (:class:`mindspore.dtype`, optional): defaults to None. Overrides the dtype of the
819            output Tensor.
820
821    Returns:
822        Tensor, sum_along_diagonals.
823
824    Raises:
825        ValueError: if the input tensor has less than two dimensions.
826
827    Supported Platforms:
828        ``Ascend`` ``GPU`` ``CPU``
829
830    Examples:
831        >>> import mindspore.numpy as np
832        >>> x = np.eye(3)
833        >>> print(x.trace())
834        3.0
835    """
836    d = x.diagonal(offset, axis1=axis1, axis2=axis2)
837    shape = d.shape
838    if dtype is None:
839        dtype = d.dtype
840    if shape[-1] == 0:
841        return F.fill(dtype, shape[:-1], 0)
842    res = F.reduce_sum(d.astype(mstype.float32), -1)
843    return res.astype(dtype)
844
845
846def take(x, indices, axis=None, mode='clip'):
847    """
848    Takes elements from an array along an axis.
849
850    Args:
851        a (Tensor): Source array with shape `(Ni…, M, Nk…)`.
852        indices (Tensor): The indices with shape `(Nj...)` of the values to extract.
853        axis (int, optional): The axis over which to select values. By default,
854            the flattened input array is used.
855        mode (‘raise’, ‘wrap’, ‘clip’, optional):
856
857            - edge: Pads with the edge values of `arr`.
858            - raise: Raises an error;
859            - wrap: Wraps around;
860            - clip: Clips to the range. `clip` mode means that all indices that are
861              too large are replaced by the index that addresses the last element
862              along that axis. Note that this disables indexing with negative numbers.
863
864    Returns:
865        Tensor, the indexed result.
866
867    Raises:
868        ValueError: if axis is out of range.
869        TypeError: if the input is not a Tensor.
870
871    Supported Platforms:
872        ``Ascend`` ``GPU`` ``CPU``
873
874    Examples:
875        >>> import mindspore.numpy as np
876        >>> a = np.array([4, 3, 5, 7, 6, 8])
877        >>> indices = np.array([0, 1, 4])
878        >>> output = a.take(indices)
879        >>> print(output)
880        [4 3 6]
881    """
882    if mode not in ('raise', 'wrap', 'clip'):
883        const_utils.raise_value_error('raise should be one of "raise", "wrap", or "clip"')
884    if axis is None:
885        a = x.ravel()
886        axis = 0
887    else:
888        a = x
889    ndim = a.ndim
890    axis = check_axis_in_range_const(axis, ndim)
891
892    shape_a = a.shape
893    shape_indices = indices.shape
894    size_indices = indices.size
895    indices = compile_utils.check_indices(shape_a[axis], indices, mode)
896
897    # reshapes indices to shape (Ni..., Nj..., Nk)
898    shape_ni = tuple_slice(shape_a, None, axis)
899    shape_nk = tuple_slice(shape_a, axis + 1, None)
900    shape_out = shape_ni + shape_indices + shape_nk
901    shape_indices = expanded_shape(ndim, size_indices, axis)
902    indices = indices.reshape(shape_indices)
903    shape_indices = shape_ni + (indices.size,) + shape_nk
904    indices = P.BroadcastTo(shape_indices)(indices)
905
906    res = F.gather_d(a, axis, indices)
907    return res.reshape(shape_out)
908
909
910def choose(x, choices, mode='clip'):
911    """
912    Construct an array from an index array and a list of arrays to choose from.
913
914    Args:
915        choices (sequence of arrays): Choice arrays. `a` and all of the `choices` must
916            be broadcastable to the same shape. If `choices` is itself an array, then
917            its outermost dimension (i.e., the one corresponding to ``choices.shape[0]``)
918            is taken as defining the “sequence”.
919        mode (‘raise’, ‘wrap’, ‘clip’, optional): Specifies how indices outside
920            ``[0, n-1]`` will be treated:
921
922            ‘raise’ – raise an error (default);
923
924            ‘wrap’ – wrap around;
925
926            ‘clip’ – clip to the range. ‘clip’ mode means that all indices that are
927            too large are replaced by the index that addresses the last element
928            along that axis. Note that this disables indexing with negative numbers.
929
930    Returns:
931        Tensor, the merged result.
932
933    Supported Platforms:
934        ``Ascend`` ``GPU`` ``CPU``
935
936    Raises:
937        ValueError: if ``len(condlist) != len(choicelist)``.
938
939    Examples:
940        >>> import mindspore.numpy as np
941        >>> choices = [[0, 1, 2, 3], [10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33]]
942        >>> x = np.array([2, 3, 1, 0])
943        >>> print(x.choose(choices))
944        [20 31 12  3]
945    """
946    if check_is_tensor(F.typeof(choices)):
947        shape_choice = infer_out_shape(x.shape, choices.shape[1:])
948        choices = P.BroadcastTo((choices.shape[0],) + shape_choice)(choices)
949    else:
950        # broadcasts choices to the same shape if choices is a sequence
951        choicelist = []
952        shapes = ()
953        for choice in choices:
954            if not check_is_tensor(F.typeof(choice)):
955                choice = const_utils.make_tensor(choice)
956            shapes += (choice.shape,)
957            choicelist.append(choice)
958        shape_choice = infer_out_shape(x.shape, *shapes)
959        tmp = []
960        for choice in choicelist:
961            tmp.append(P.BroadcastTo(shape_choice)(choice))
962        choices = F.stack(tmp)
963
964    if x.ndim == 0 or choices.ndim == 0:
965        const_utils.raise_value_error('input cannot be scalars')
966    a = P.BroadcastTo(shape_choice)(x)
967    dtype = choices.dtype
968    # adjusts dtype for F.tensor_mul and F.gather_nd
969    a = a.astype(mstype.int32)
970    choices = choices.astype(mstype.int32)
971    a = compile_utils.check_indices(choices.shape[0], a, mode, allow_negative_index=False)
972
973    grids = []
974    ndim = len(a.shape)
975    for i in range(ndim):
976        dim_grid = const_utils.make_tensor(F.make_range(a.shape[i]), mstype.int32)
977        dim_shape = expanded_shape(ndim, a.shape[i], i)
978        dim_grid = P.BroadcastTo(a.shape)(dim_grid.reshape(dim_shape))
979        grids.append(dim_grid)
980    grid = P.Stack(-1)(grids)
981    indices = P.Concat(-1)((a.reshape(a.shape + (1,)), grid))
982    return F.gather_nd(choices, indices).astype(dtype)
983
984
985def searchsorted(x, v, side='left', sorter=None):
986    """
987    Finds indices where elements should be inserted to maintain order.
988
989    Args:
990        v (Union[int, float, bool, list, tuple, Tensor]): Values to insert into `a`.
991        side ('left', 'right', optional): If ‘left’, the index of the first suitable
992            location found is given. If ‘right’, return the last such index. If there is
993            no suitable index, return either 0 or N (where N is the length of `a`).
994        sorter (Union[int, float, bool, list, tuple, Tensor]): 1-D optional array of
995            integer indices that sort array `a` into ascending order. They are typically
996            the result of argsort.
997
998    Returns:
999        Tensor, array of insertion points with the same shape as `v`.
1000
1001    Raises:
1002        ValueError: if argument for `side` or `sorter` is invalid.
1003
1004    Supported Platforms:
1005        ``Ascend`` ``GPU`` ``CPU``
1006
1007    Examples:
1008        >>> from mindspore import numpy as np
1009        >>> x = np.array([1,2,3,4,5])
1010        >>> print(x.searchsorted(3))
1011        2
1012    """
1013    if side not in ('left', 'right'):
1014        const_utils.raise_value_error('invalid value for keyword "side"')
1015    a = x.astype(mstype.float32)
1016    if not check_is_tensor(F.typeof(v)):
1017        v = const_utils.make_tensor(v)
1018    shape = v.shape
1019    if sorter is not None:
1020        if sorter.ndim != 1 or sorter.size != a.size:
1021            const_utils.raise_value_error('sorter must be 1-D array with the same size as `a`')
1022        sorter = const_utils.make_tensor(sorter)
1023        sorter = sorter.reshape(sorter.shape + (1,))
1024        a = F.gather_nd(a, sorter)
1025    less_op = F.tensor_le if side == 'left' else F.tensor_lt
1026    i = F.fill(mstype.int32, shape, 0)
1027    j = F.fill(mstype.int32, shape, a.size)
1028
1029    sort_range = F.make_range(get_log2_size(F.shape_mul(shape) + 1))
1030    for _ in sort_range:
1031        mid = (i - F.neg_tensor(j))//2
1032        mask = less_op(v, F.gather_nd(a, mid.reshape(mid.shape + (1,))))
1033        i = F.select(mask, i, mid)
1034        j = F.select(mask, mid, j)
1035    return j
1036
1037
1038def fill(x, value):
1039    """
1040    Fills the array with a scalar value.
1041
1042    Note:
1043        Unlike Numpy, tensor.fill() will always returns a new tensor, instead of
1044        filling the original tensor.
1045
1046    Args:
1047        value (Union[None, int, float, bool]): All elements of a will be assigned this value.
1048
1049    Returns:
1050        Tensor, with the original dtype and shape as input tensor.
1051
1052    Raises:
1053        TypeError: If input arguments have types not specified above.
1054        ValueError: If `shape` has entries < 0.
1055
1056    Supported Platforms:
1057        ``Ascend`` ``GPU`` ``CPU``
1058
1059    Examples:
1060        >>> import numpy as np
1061        >>> from mindspore import Tensor
1062        >>> a = Tensor(np.arange(4).reshape((2,2)).astype('float32'))
1063        >>> print(a.fill(1.0))
1064        [[1. 1.]
1065        [1. 1.]]
1066    """
1067    if value is None:
1068        if x.dtype not in (mstype.float16, mstype.float32, mstype.float64):
1069            const_utils.raise_type_error("If None is used as value, the original Tensor's dtype must be float.")
1070        value = nan_tensor
1071        return F.tile(value, x.shape).astype(x.dtype)
1072    if not isinstance(value, (int, float, bool)):
1073        const_utils.raise_type_error("input value must be a scalar.")
1074    return F.fill(x.dtype, x.shape, value)
1075
1076
1077def ptp(x, axis=None, keepdims=False):
1078    """
1079    The name of the function comes from the acronym for ‘peak to peak’.
1080
1081    Note:
1082        Numpy arguments `dtype` and `out` are not supported.
1083
1084    Args:
1085        x (Tensor): Input tensor.
1086        axis (Union[None, int, tuple(int)]): Axis or axes along which the range is computed.
1087            The default is to compute the variance of the flattened array. Default: None.
1088        keepdims (bool): Default is False.
1089
1090    Returns:
1091        Tensor.
1092
1093    Raises:
1094        TypeError: if the input is not a tensor.
1095
1096    Supported Platforms:
1097        ``Ascend`` ``GPU`` ``CPU``
1098
1099    Examples:
1100        >>> from mindspore import Tensor
1101        >>> x = Tensor([[4.0, 9.0, 2.0, 10.0], [6.0, 9.0, 7.0, 12.0]]).astype("float32")
1102        >>> print(x.ptp(axis=1))
1103        [8. 6.]
1104        >>> print(x.ptp(axis=0))
1105        [2. 0. 5. 2.]
1106    """
1107    if not isinstance(keepdims, bool):
1108        const_utils.raise_type_error('keepdims should be boolean')
1109    if axis is None:
1110        axis = ()
1111    else:
1112        check_axis_type(axis, True, True, False)
1113        axis = check_axis_valid(axis, x.ndim)
1114
1115    return x.max(axis, keepdims) - x.min(axis, keepdims)
1116
1117
1118def clip(x, xmin, xmax, dtype=None):
1119    """
1120    Clips (limits) the values in an array.
1121
1122    Given an interval, values outside the interval are clipped to the interval edges.
1123    For example, if an interval of :math:`[0, 1]` is specified, values smaller than 0 become 0,
1124    and values larger than 1 become 1.
1125
1126    Note:
1127        Currently, clip with `nan` is not supported.
1128
1129    Args:
1130        x (Tensor): Tensor containing elements to clip.
1131        xmin (Tensor, scalar, None): Minimum value. If None, clipping is not performed
1132            on lower interval edge. Not more than one of `xmin` and `xmax` may be None.
1133        xmax (Tensor, scalar, None): Maximum value. If None, clipping is not performed
1134            on upper interval edge. Not more than one of `xmin` and `xmax` may be None.
1135            If `xmin` or `xmax` are tensors, then the three tensors will be broadcasted
1136            to match their shapes.
1137        dtype (:class:`mindspore.dtype`, optional): defaults to None. Overrides the dtype of the
1138            output Tensor.
1139
1140    Returns:
1141        Tensor, a tensor with the elements of `x`, but where values
1142        < `xmin` are replaced with `xmin`, and those > `xmax` with `xmax`.
1143
1144    Supported Platforms:
1145        ``Ascend`` ``GPU`` ``CPU``
1146
1147    Examples:
1148        >>> from mindspore import Tensor
1149        >>> x = Tensor([1, 2, 3, -4, 0, 3, 2, 0]).astype("float32")
1150        >>> output = x.clip(0, 2)
1151        >>> print(output)
1152        [1 2 2 0 0 2 2 0]
1153    """
1154    if xmin is None and xmax is None:
1155        const_utils.raise_value_error("One of max or min must be given.")
1156    is_scalar = False
1157    if xmin is not None:
1158        xmin = const_utils.make_tensor(xmin).astype(x.dtype)
1159        if x.ndim == 0 and xmin.ndim == 0:
1160            x = F.maximum(x.reshape((1,)), xmin).squeeze()
1161        else:
1162            x = F.maximum(x, xmin)
1163    if xmax is not None:
1164        xmax = const_utils.make_tensor(xmax).astype(x.dtype)
1165        if x.ndim == 0 and xmax.ndim == 0:
1166            x = F.minimum(x.reshape((1,)), xmax).squeeze()
1167        else:
1168            x = F.minimum(x, xmax)
1169    if is_scalar:
1170        return x.squeeze()
1171    if dtype is not None and dtype != x.dtype:
1172        return x.astype(dtype)
1173    return x
1174
1175
1176def var(x, axis=None, ddof=0, keepdims=False):
1177    """
1178    Compute the variance along the specified axis.
1179    The variance is the average of the squared deviations from the mean, i.e.,
1180    :math:`var = mean(abs(x - x.mean())**2)`.
1181
1182    Return the variance, which is computed for the flattened array by default,
1183    otherwise over the specified axis.
1184
1185    Note:
1186        Numpy arguments `dtype`, `out` and `where` are not supported.
1187
1188    Args:
1189        x (Tensor): A Tensor to be calculated.
1190        axis (Union[None, int, tuple(int)]): Axis or axes along which the variance is computed.
1191            The default is to compute the variance of the flattened array. Default: `None`.
1192        ddof (int): Means Delta Degrees of Freedom. Default: 0.
1193            The divisor used in calculations is :math:`N - ddof`, where :math:`N` represents the number of elements.
1194        keepdims (bool): Default: `False`.
1195
1196    Supported Platforms:
1197        ``Ascend`` ``GPU`` ``CPU``
1198
1199    Returns:
1200        Standard deviation tensor.
1201
1202    Examples:
1203        >>> import mindspore.numpy as np
1204        >>> input_x = np.array([1., 2., 3., 4.])
1205        >>> print(input_x.var())
1206        1.25
1207    """
1208    if 0 in x.shape:
1209        return nan_tensor.astype(x.dtype)
1210    if not isinstance(ddof, int) or not isinstance(keepdims, int):
1211        const_utils.raise_type_error("integer argument expected")
1212
1213    if axis is None:
1214        axis = ()
1215    else:
1216        axis = check_and_canonicalize_axes(axis, x.ndim)
1217    x_mean = _mean_keepdims(x, axis)
1218    x_sub = F.tensor_sub(x, x_mean)
1219    x_pow = F.tensor_pow(x_sub, 2)
1220    if keepdims:
1221        x_sum = _reduce_sum_keepdims(x_pow, axis)
1222    else:
1223        x_sum = _reduce_sum_default(x_pow, axis)
1224
1225    if axis == ():
1226        axis = F.make_range(x.ndim)
1227    nums = 1
1228    for ax in axis:
1229        nums *= x.shape[ax]
1230    return F.tensor_div(x_sum, nums - ddof)
1231
1232
1233def std(x, axis=None, ddof=0, keepdims=False):
1234    """
1235    Compute the standard deviation along the specified axis.
1236    The standard deviation is the square root of the average of the squared deviations
1237    from the mean, i.e., :math:`std = sqrt(mean(abs(x - x.mean())**2))`.
1238
1239    Return the standard deviation, which is computed for the flattened array by default,
1240    otherwise over the specified axis.
1241
1242    Note:
1243        Numpy arguments `dtype`, `out` and `where` are not supported.
1244
1245    Args:
1246        x (Tensor): A Tensor to be calculated.
1247        axis (Union[None, int, tuple(int)]): Axis or axes along which the standard
1248            deviation is computed. Default: `None`.
1249
1250            If `None`, compute the standard deviation of the flattened array.
1251        ddof (int): Means Delta Degrees of Freedom. The divisor used in calculations is :math:`N - ddof`,
1252            where :math:`N` represents the number of elements. Default: 0.
1253        keepdims: Default: `False`.
1254
1255    Returns:
1256        Standard deviation tensor.
1257
1258    Supported Platforms:
1259        ``Ascend`` ``GPU`` ``CPU``
1260
1261    Examples:
1262        >>> import mindspore.numpy as np
1263        >>> input_x = np.array([1., 2., 3., 4.])
1264        >>> print(input_x.std())
1265        1.118034
1266    """
1267    x_var = var(x, axis, ddof, keepdims)
1268    return F.tensor_pow(x_var, 0.5)
1269
1270
1271def sum(x, axis=None, dtype=None, keepdims=False, initial=None): # pylint: disable=redefined-builtin
1272    """
1273    Return sum of array elements over a given axis.
1274
1275    Note:
1276        Numpy arguments `out`, `where`, `casting`, `order`, `subok`, `signature`, and
1277        `extobj` are not supported.
1278
1279    Args:
1280        x (Union[int, float, bool, list, tuple, Tensor]): Elements to sum.
1281        axis (Union[None, int, tuple(int)]): Axis or axes along which a sum is performed. Default: None.
1282            If None, sum all of the elements of the input array.
1283            If axis is negative it counts from the last to the first axis.
1284            If axis is a tuple of ints, a sum is performed on all of the axes specified in the tuple
1285            instead of a single axis or all the axes as before.
1286        dtype (:class:`mindspore.dtype`, optional): defaults to None. Overrides the dtype of the
1287            output Tensor.
1288        keepdims (bool): If this is set to True, the axes which are reduced are left in the result as
1289            dimensions with size one. With this option, the result will broadcast correctly against the input array.
1290            If the default value is passed, then keepdims will not be passed through to the sum method of
1291            sub-classes of ndarray, however any non-default value will be. If the sub-class’ method does not
1292            implement keepdims any exceptions will be raised.
1293        initial (scalar): Starting value for the sum.
1294
1295    Returns:
1296        Tensor. A tensor with the same shape as input, with the specified axis removed.
1297        If input tensor is a 0-d array, or if axis is None, a scalar is returned.
1298
1299    Raises:
1300        TypeError: If input is not array_like or `axis` is not int or tuple of ints or
1301            `keepdims` is not integer or `initial` is not scalar.
1302        ValueError: If any axis is out of range or duplicate axes exist.
1303
1304    Supported Platforms:
1305        ``Ascend`` ``GPU`` ``CPU``
1306
1307    Examples:
1308        >>> import mindspore.numpy as np
1309        >>> input_x = np.array([-1, 0, 1]).astype('int32')
1310        >>> print(input_x.sum())
1311        0
1312        >>> input_x = np.arange(10).reshape(2, 5).astype('float32')
1313        >>> print(input_x.sum(axis=1))
1314        [10. 35.]
1315    """
1316    input_x = x.astype(mstype.int32) if x.dtype == mstype.bool_ else x
1317    dtype = input_x.dtype if dtype is None else dtype
1318    if not isinstance(keepdims, int):
1319        const_utils.raise_type_error("integer argument expected")
1320    if initial is not None and not isinstance(initial, (int, float, bool)):
1321        const_utils.raise_type_error("initial argument should be a scalar.")
1322    if axis is None:
1323        axis = ()
1324    else:
1325        axis = check_and_canonicalize_axes(axis, x.ndim)
1326
1327    if not check_type_support(input_x.dtype, 'GPU', (mstype.float64, mstype.float32, mstype.float16)):
1328        input_x = input_x.astype(mstype.float32)
1329    if 0 in x.shape:
1330        x = const_utils.make_tensor([0], x.dtype)
1331    if keepdims:
1332        res = _reduce_sum_keepdims(input_x, axis)
1333    else:
1334        res = _reduce_sum_default(input_x, axis)
1335    if initial is not None:
1336        res += initial
1337    return res.astype(dtype)
1338
1339
1340def repeat(x, repeats, axis=None):
1341    """
1342    Repeat elements of an array.
1343
1344    Args:
1345        x (Tensor): Input tensor.
1346        repeats (Union[int, tuple, list]): The number of repetitions for each element.
1347            `repeats` is broadcasted to fit the shape of the given axis.
1348        axis (int, optional): The axis along which to repeat values. By default,
1349            use the flattened input tensor, and return a flat output tensor.
1350
1351    Returns:
1352        Tensor, has the same shape as input tensor except along the given axis.
1353
1354    Raises:
1355        ValueError: if axis is out of range.
1356        TypeError: if input is not a Tensor.
1357
1358    Supported Platforms:
1359        ``Ascend`` ``GPU`` ``CPU``
1360
1361    Examples:
1362        >>> import mindspore.numpy as np
1363        >>> x = np.array(3)
1364        >>> print(x.repeat(4))
1365        [3 3 3 3]
1366        >>> x = np.array([[1,2],[3,4]])
1367        >>> print(x.repeat(2))
1368        [1 1 2 2 3 3 4 4]
1369        >>> print(x.repeat(3, axis=1))
1370        [[1 1 1 2 2 2]
1371        [3 3 3 4 4 4]]
1372        >>> print(x.repeat([1,2], axis=0))
1373        [[1 2]
1374        [3 4]
1375        [3 4]]
1376    """
1377    if not isinstance(repeats, (tuple, list)):
1378        repeats = (repeats,)
1379    for element in repeats:
1380        if not isinstance(element, int):
1381            const_utils.raise_type_error("Each element should be integer")
1382    if axis is None:
1383        x = ravel(x)
1384        axis = 0
1385    if not isinstance(axis, int):
1386        const_utils.raise_type_error('axes should be integers')
1387    check_axis_in_range_const(axis, x.ndim)
1388    axis = axis + x.ndim if axis < 0 else axis
1389
1390    if len(repeats) == 1:
1391        repeats = repeats[0]
1392        if repeats == 0:
1393            return empty_tensor(x.dtype)
1394        return repeat_elements(x, repeats, axis)
1395    size = x.shape[axis]
1396    if len(repeats) != size:
1397        const_utils.raise_value_error('operands could not be broadcast together')
1398    subs = P.Split(axis, size)(x)
1399    repeated_subs = []
1400    for sub, rep in zip(subs, repeats):
1401        if rep != 0:
1402            repeated_subs.append(repeat_elements(sub, rep, axis))
1403    return P.Concat(axis)(repeated_subs)
1404
1405
1406def getitem(data, index):
1407    """Implementation of `getitem`."""
1408    return data.__getitem__(index)
1409
1410
1411def setitem(data, index, value):
1412    """Implementation of `setitem`."""
1413    return data.__setitem__(index, value)
1414
1415
1416def item(data, *args):
1417    """Implementation of `item`."""
1418    return compile_utils.tensor_item(data, *args)
1419
1420
1421def itemset(data, *args):
1422    """Implementation of `itemset`."""
1423    return compile_utils.tensor_itemset(data, *args)
1424
1425
1426def ms_iter(xs):
1427    """Implementation of `iter`."""
1428    return xs.__ms_iter__()
1429
1430
1431def ms_next(it):
1432    """Implementation of `next`."""
1433    return it.__ms_next__()
1434
1435
1436def hasnext(it):
1437    """Implementation of `hasnext`."""
1438    return it.__ms_hasnext__()
1439
1440
1441def ms_len(data):
1442    """Implementation of `len`."""
1443    return data.__len__()
1444
1445
1446def floor(x):
1447    """Implementation of `floor`."""
1448    return x.__floor__()
1449
1450
1451def trunc(x):
1452    """Implementation of `trunc`."""
1453    return x.__trunc__()
1454
1455
1456def uadd(x):
1457    """Implementation of `uadd`."""
1458    return x.__pos__()
1459
1460
1461def usub(x):
1462    """Implementation of `usub`."""
1463    return x.__neg__()
1464
1465
1466def scalar_truediv(x, y):
1467    """Implementation of `scalar_truediv`."""
1468    return x.__truediv__(y)
1469
1470
1471def scalar_floordiv(x, y):
1472    """Implementation of `scalar_floordiv`."""
1473    return x.__floordiv__(y)
1474
1475
1476def bool_(x):
1477    """Implementation of `bool`."""
1478    return x.__bool__()
1479
1480
1481def enumerate_(x, start=0):
1482    """Enumerate list or tuple or tensor."""
1483    x_type = F.typeof(x)
1484    ret = ()
1485    op_name = "enumerate"
1486    if check_is_tuple_or_list_or_tensor(x_type, op_name, "first input") and \
1487            check_is_const_int(start, op_name, "start"):
1488        if check_is_tensor(x_type):
1489            for i in range(x.shape[0]):
1490                ret += ((start + i, x[i]),)
1491        else:
1492            ret = zip(range(start, start + len(x)), x)
1493    return ret
1494
1495
1496def expand_tensor_as(x, y):
1497    """Expand tensor"""
1498    broadcast_to = P.BroadcastTo(shape_(y))
1499    return broadcast_to(x)
1500
1501
1502def view(x, *shape):
1503    """Reshape tensor, if shape is -1, reshape tensor into one dimension"""
1504    shape = check_view_shape(shape)
1505    return F.reshape(x, shape)
1506
1507
1508def isinstance_(x, base_type):
1509    """Determine whether x is an instance of base_type."""
1510    x_type = F.typeof(x)
1511    return check_type_same(x_type, base_type)
1512
1513
1514def while_cond(x):
1515    """For while condition, if the condition is a tensor, the loop will not be unrolled"""
1516    if F.issubclass_(F.typeof(x), F.typeof(mstype.tensor)):
1517        is_cond = check_is_tensor_bool_cond(F.shape(x))
1518        if is_cond:
1519            return F.cast(x, mstype.bool_)
1520    return x
1521
1522
1523@constexpr
1524def empty_tensor(dtype):
1525    return Tensor([], dtype)
1526
1527
1528@constexpr
1529def check_type_same(x_type, base_type):
1530    """Check x_type is same as base_type."""
1531    pytype_to_mstype = {
1532        bool: mstype.Bool,
1533        int: mstype.Int,
1534        float: mstype.Float,
1535        str: mstype.String,
1536        list: mstype.List,
1537        tuple: mstype.Tuple,
1538        dict: mstype.Dict,
1539        Tensor: mstype.tensor_type,
1540        Parameter: mstype.ref_type
1541    }
1542
1543    has_int = False
1544    has_tensor = False
1545
1546    def to_target_type(origin_type):
1547        try:
1548            if isinstance(origin_type, type):
1549                ret_type = pytype_to_mstype[origin_type]
1550                if ret_type == mstype.Int:
1551                    nonlocal has_int
1552                    has_int = True
1553                if ret_type == mstype.tensor_type:
1554                    nonlocal has_tensor
1555                    has_tensor = True
1556                return (ret_type,)
1557            if isinstance(origin_type, tuple):
1558                return tuple(to_target_type(i) for i in origin_type)
1559            raise TypeError(f"The second arg of 'isinstance' must be a type or a tuple of types, "
1560                            f"but got a {type(origin_type).__name__}")
1561        except KeyError:
1562            raise TypeError(f"The second arg of 'isinstance' should be bool, int, float, str, list, tuple, "
1563                            f"Tensor, Parameter, or a tuple containing only these types, but got {origin_type}")
1564    target_type = to_target_type(base_type)
1565    if (isinstance(x_type, mstype.Bool) and has_int) or (isinstance(x_type, mstype.ref_type) and has_tensor):
1566        return True
1567    return isinstance(x_type, target_type)
1568
1569
1570@constexpr
1571def get_itemsize(x_type):
1572    """get itemsize from tensor's dtype."""
1573    return itemsize_map[x_type]
1574
1575
1576@constexpr
1577def check_is_tensor(x):
1578    """check whether x is tensor."""
1579    if isinstance(x, mstype.tensor_type):
1580        return True
1581    return False
1582
1583
1584@constexpr
1585def check_is_tuple_or_list_or_tensor(x, op_name, arg_name):
1586    """check whether x is list or tuple or tensor."""
1587    if isinstance(x, (mstype.List, mstype.Tuple, mstype.tensor_type)):
1588        return True
1589    raise TypeError(f"For '{op_name}', the '{arg_name}' should be tuple or list or tensor, but got {x}.")
1590
1591
1592@constexpr
1593def check_is_const_int(x, op_name, arg_name):
1594    """check whether x is const int."""
1595    if x is None:
1596        raise TypeError(f"For '{op_name}', the '{arg_name}' should be a const int number, but got not const.")
1597    if not isinstance(x, int):
1598        raise TypeError(f"For '{op_name}', the '{arg_name}' should be a const int number, but got {x}.")
1599    return True
1600
1601
1602@constexpr
1603def check_is_tensor_bool_cond(shp):
1604    """check if tensor is a bool condition"""
1605    if shp in ((), (1,)):
1606        return True
1607    raise ValueError(f"Only tensor which shape is () or (1,) can be converted to bool, but got tensor shape is {shp}")
1608
1609
1610@constexpr
1611def const_tensor_to_bool(x):
1612    """convert bool tensor to bool condition"""
1613    if x is None:
1614        raise ValueError("Only tensor which shape is () or (1,) can be converted to bool, but got None")
1615    x = x.asnumpy()
1616    if x.shape == ():
1617        return bool(x)
1618    if x.shape == (1,):
1619        return bool(x[0])
1620    raise ValueError(
1621        f"Only tensor which shape is () or (1,) can be converted to bool, but got tensor shape is {x.shape}")
1622
1623
1624@constexpr
1625def check_view_shape(x):
1626    """Check view function input shape"""
1627    if not x:
1628        raise ValueError("The shape variable should not be empty")
1629    if isinstance(x[0], tuple):
1630        if len(x) != 1:
1631            raise ValueError(f"Only one tuple is needed, but got {x}")
1632        x = x[0]
1633    return x
1634
1635
1636# convert normal param_check functions to constexpr functions
1637check_astype_dtype_const = constexpr(validator.check_astype_dtype)
1638check_transpose_axis_const = constexpr(validator.check_transpose_axis)
1639check_reshape_shp_const = constexpr(validator.check_reshape_shp)
1640check_flatten_order_const = constexpr(validator.check_flatten_order)
1641check_swapaxes_axis_const = constexpr(validator.check_swapaxes_axis)
1642prepare_shape_for_squeeze_const = constexpr(validator.prepare_shape_for_squeeze)
1643check_axis_in_range_const = constexpr(validator.check_axis_in_range)
1644check_axis_valid = constexpr(validator.check_axis_valid)
1645max_ = constexpr(validator.max_)
1646min_ = constexpr(validator.min_)
1647expanded_shape = constexpr(validator.expanded_shape)
1648tuple_slice = constexpr(validator.tuple_slice)
1649infer_out_shape = constexpr(validator.infer_out_shape)
1650get_log2_size = constexpr(validator.get_log2_size)
1651check_axis_type = constexpr(validator.check_axis_type)
1652check_and_canonicalize_axes = constexpr(validator.check_and_canonicalize_axes)
1653empty_compile = constexpr(validator.empty_compile)
1654check_type_support = constexpr(validator.check_type_support)
1655
1656
1657def tensor_bool(x):
1658    """tensor as condition, if is constant, return immediate bool value"""
1659    is_cond = check_is_tensor_bool_cond(F.shape(x))
1660    if is_cond and F.isconstant(x):
1661        return const_tensor_to_bool(x)
1662    return F.cast(x, mstype.bool_)
1663
1664
1665def and_(x, y):
1666    """Implementation of `and` (`&`)."""
1667    return x.__and__(y)
1668
1669
1670def or_(x, y):
1671    """Implementation of `or` (`|`)."""
1672    return x.__or__(y)
1673
1674
1675def matmul(x, y):
1676    """Implementation of `matmul` (`@`)."""
1677    return x.__matmul__(y)
1678
1679
1680def float_bool(x):
1681    """Implementation of `float_bool`."""
1682    return x != 0.0
1683
1684
1685def int_bool(x):
1686    """Implementation of `int_bool`."""
1687    return x != 0
1688
1689
1690def str_bool(x):
1691    """Implementation of `str_bool`."""
1692    if x == "":
1693        return False
1694    return True
1695
1696
1697def list_bool(x):
1698    """Implementation of `tuple_bool`."""
1699    return len(x) != 0
1700
1701
1702def tuple_bool(x):
1703    """Implementation of `tuple_bool`."""
1704    return len(x) != 0
1705
1706
1707def dict_bool(x):
1708    """Implementation of `dict_bool`."""
1709    return len(x) != 0
1710
1711
1712def none_bool(x):
1713    """Implementation of `none_bool`."""
1714    return False
1715
1716
1717def func_bool(x):
1718    """Implementation of `func_bool`."""
1719    return True
1720
1721
1722def float_floordiv(x, y):
1723    """Implementation of `float_floordiv`."""
1724    return floor(x / y)
1725
1726
1727#############
1728# Iteration #
1729#############
1730
1731
1732@dataclass(frozen=True)
1733class SequenceIterator:
1734    """
1735    SequenceIterator is a util dataclass for iterating sequence object.
1736
1737    Iterator to use for sequences like List, Array.
1738    """
1739
1740    idx: int
1741    seq: list
1742
1743    @core(ignore_values=True)
1744    def __ms_hasnext__(self):
1745        """Whether the index is past the length of the sequence."""
1746        return self.idx < ms_len(self.seq)
1747
1748    @core(ignore_values=True)
1749    def __ms_next__(self):
1750        """Return the next element and a new iterator."""
1751        return self.seq[self.idx], SequenceIterator(self.idx + 1, self.seq)
1752
1753
1754def list_iter(xs):
1755    """Iterator for List."""
1756    return SequenceIterator(0, xs)
1757
1758
1759def array_iter(xs):
1760    """Iterator for Array."""
1761    return SequenceIterator(0, xs)
1762
1763
1764def tuple_next(xs):
1765    """Next tuple."""
1766    return xs[0], tail(xs)
1767
1768
1769def tuple_hasnext(xs):
1770    """Whether the tuple is empty or not."""
1771    return len(xs) > 0
1772
1773
1774def list_next(xs):
1775    """Next list."""
1776    return xs[0], tail(xs)
1777
1778
1779def list_hasnext(xs):
1780    """Whether the list is empty or not."""
1781    return len(xs) > 0
1782
1783
1784# pylint: disable=redefined-outer-name
1785def list_append(self_, item):
1786    return _append(self_, item)
1787
1788
1789#################
1790# Array methods #
1791#################
1792
1793
1794def to_array(x):
1795    """Implementation of `to_array`."""
1796    return x.__ms_to_array__()
1797