cunumeric.tensordot#

cunumeric.tensordot(a: ndarray, b: ndarray, axes: int | tuple[int | Sequence[int], int | Sequence[int]] = 2, out: ndarray | None = None) → ndarray#

Compute 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.

Parameters:

a (array_like) – Tensors to “dot”.
b (array_like) – Tensors to “dot”.
axes (int or array_like) –
- integer_like If an int N, sum over the last N axes of a and the first N axes of b in order.
- (2,) array_like 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.
out (ndarray, optional) – Output argument. This must have the exact shape that would be returned if it was not present. If its dtype is not what would be expected from this operation, then the result will be (unsafely) cast to out.

Returns:

output – The tensor dot product of the inputs. If out is given, then it is returned.

Return type:

ndarray

Notes

The cuNumeric implementation is a little more liberal than NumPy in terms of allowed broadcasting, e.g. tensordot(ones((3,1)), ones((1,4))) is allowed.

Except for the inner-product case, only floating-point types are supported.