mindspore.numpy.tensordot

mindspore.numpy.tensordot(a, b, axes=2)[源代码]

Computes tensor dot product along specified axes.

Given two tensors, a and b, and an array_like object containing two array_like objects, (a_axes, b_axes), sum the products of a’s and b’s elements (components) over the axes specified by a_axes and b_axes. The third argument can be a single non-negative integer_like scalar, N; if it is such, then the last N dimensions of a and the first N dimensions of b are summed over. Three common use cases are:

  • axes = 0 : tensor product

  • axes = 1 : tensor dot product

  • axes = 2 : (default) tensor double contraction

When axes is integer_like, the sequence for evaluation will be: first the -Nth axis in a and 0th axis in b, and the -1th axis in a and Nth axis in b last. When there is more than one axis to sum over - and they are not the last (first) axes of a (b) - the argument axes should consist of two sequences of the same length, with the first axis to sum over given first in both sequences, the second axis second, and so forth. The shape of the result consists of the non-contracted axes of the first tensor, followed by the non-contracted axes of the second.

Note

On CPU, the supported dypes are np.float16 and np.float32. On GPU, the supported dypes are np.float16 and np.float32.

Parameters

  • a (Tensor) – Tensor to “dot”.

  • b (Tensor) – Tensor to “dot”.

  • axes (int or sequence of ints) –

    integer_like: If an int N, sum over the last N axes of a and the first N axes of b in order. The sizes of the corresponding axes must match.

    sequence of ints: Or, a list of axes to be summed over, first sequence applying to a, second to b. Both elements array_like must be of the same length.

Returns

Tensor, or list of tensors, the tensor dot product of the input.

Supported Platforms:

Ascend GPU CPU

Examples

>>> import mindspore.numpy as np
>>> a = np.ones((3, 4, 5))
>>> b = np.ones((4, 3, 2))
>>> output = np.tensordot(a, b, axes=([1,0],[0,1]))
>>> print(output.shape)
(5, 2)