tensordot — Python array API standard 2024.12 documentation (original) (raw)

tensordot(x1: array, x2: array, /, *, axes: int | Tuple[Sequence[int], Sequence[int]] = 2) → array¶

Returns a tensor contraction of x1 and x2 over specific axes.

Note

The tensordot function corresponds to the generalized matrix product.

Parameters:

• x1 (array) – first input array. Should have a numeric data type.
• x2 (array) –
  second input array. Should have a numeric data type. Corresponding contracted axes of x1 and x2 must be equal.
  Note
  Contracted axes (dimensions) must not be broadcasted.
• axes (Union [ int , Tuple [ Sequence [ int ] , Sequence [ int ] ] ]) –
  number of axes (dimensions) to contract or explicit sequences of axis (dimension) indices for x1 and x2, respectively.
  If axes is an int equal to N, then contraction must be performed over the last N axes of x1 and the first N axes of x2 in order. The size of each corresponding axis (dimension) must match. Must be nonnegative.
  • If N equals 0, the result is the tensor (outer) product.
  • If N equals 1, the result is the tensor dot product.
  • If N equals 2, the result is the tensor double contraction (default).
    If axes is a tuple of two sequences (x1_axes, x2_axes), the first sequence must apply to x1 and the second sequence to x2. Both sequences must have the same length. Each axis (dimension) x1_axes[i] for x1 must have the same size as the respective axis (dimension) x2_axes[i] for x2. Each index referred to in a sequence must be unique. If x1 has rank (i.e, number of dimensions) N, a valid x1 axis must reside on the half-open interval [-N, N). If x2 has rank M, a valid x2 axis must reside on the half-open interval [-M, M).

Note

If either x1 or x2 has a complex floating-point data type, neither argument must be complex-conjugated or transposed. If conjugation and/or transposition is desired, these operations should be explicitly performed prior to computing the generalized matrix product.

Returns:

out (array) – an array containing the tensor contraction whose shape consists of the non-contracted axes (dimensions) of the first array x1, followed by the non-contracted axes (dimensions) of the second array x2. The returned array must have a data type determined by Type Promotion Rules.

Notes

Changed in version 2022.12: Added complex data type support.

Changed in version 2023.12: Allow negative axes.