oneflow.tensordot

oneflow.tensordot(a, b, dims=Union[int, Tensor, Tuple[List[int], List[int]], List[List[int]]], out=None)Tensor

Compute tensor dot along given dimensions.

Given two tensors a and b, and dims which represent two lists containing dim indices, tensordot traverses the two lists and calculate the tensor dot along every dim pair.

Parameters
  • a (oneflow.Tensor) – The input tensor to compute tensordot

  • b (oneflow.Tensor) – The input tensor to compute tensordot

  • dims (int or list or tuple or oneflow.Tensor) – The dims to calculate tensordot. If it’s an integer or oneflow.Tensor with only one element, the last dims of tensor a and the first dims of tensor b will be calculated. If it’s a list or tuple or oneflow.Tensor with more than one element, it must contain two array-like object, which represent the dims of tensor a and tensor b to be calculated.

  • out (oneflow.Tensor) – The tensor to save result (NOT IMPLEMENTED YET)

Returns

The result tensor

Return type

oneflow.Tensor

For example:

>>> import oneflow as flow
>>> a = flow.randn(3, 4, 5)
>>> b = flow.randn(4, 5, 6)
>>> flow.tensordot(a, b, dims=2).shape
oneflow.Size([3, 6])
>>> b = flow.randn(5, 6, 7)
>>> flow.tensordot(a, b, dims=1).shape
oneflow.Size([3, 4, 6, 7])
>>> b = flow.randn(3, 4, 7)
>>> flow.tensordot(a, b, dims=[[0, 1], [0, 1]]).shape
oneflow.Size([5, 7])

Note

Three common use cases are:

  • dims = 0 : tensor product \(a \otimes b\)

  • dims = 1 : tensor dot product \(a \cdot b\)

  • dims = 2 : (default) tensor double contraction \(a : b\)

The part of documentation is referenced from https://numpy.org/doc/stable/reference/generated/numpy.tensordot.html.

Note

The operation is equivalent to the series of operations:

  • Permute the dimensions of the tensor A that require tensordot to the end

  • Permute the dimensions of the tensor B that require tensordot to the start

  • Reshape the permuted tensor A into a 2-dimensional tensor, where the size of the 0th dimension is the product of the dimensions that do not require dot product, and the size of the 1st dimension is the product of the dimensions that require dot product

  • Reshape the permuted tensor B into a 2-dimensional tensor, where the size of the 0th dimension is the product of the dimensions that require dot product, and the size of the 1st dimension is the product of the dimensions that do not require dot product

  • Calculate the matrix multiplication of reshaped tensor A and reshaped tensor B

  • Reshape the result of matrix multiplication, the target shape is the concatenation of the dimensions that do not require tensordot of tensor A and B

This series of operations can be equivalently represented by the following code:

>>> import oneflow as flow
>>> a = flow.randn(2, 4, 3)
>>> b = flow.randn(3, 4, 2)
>>> dims = [[0, 2], [2, 0]]
>>> permuted_a = a.permute(1, 0, 2) # 0, 2 are the dimensions requiring tensordot and are placed in the end in permuting
>>> permuted_b = b.permute(2, 0, 1) # 2, 0 are the dimensions requiring tensordot and are placed at the beginning in permuting
>>> reshaped_a = permuted_a.reshape(4, 2 * 3) # 4 is the dimensions of a that do not require tensordot
>>> reshaped_b = permuted_b.reshape(2 * 3, 4) # 4 is the dimensions of a that do not require tensordot
>>> matmul_result = flow.matmul(reshaped_a, reshaped_b)
>>> result = matmul_result.reshape(4, 4) # 4, 4 are the concatentation of dimensions that do not require tensordot of a and b
>>> flow.all(result == flow.tensordot(a, b, dims))
tensor(True, dtype=oneflow.bool)