Compute the tensor dot product in Python

The tensor dot product is a generalization of matrix multiplication that allows you to sum products of tensor elements over specified axes. NumPy's tensordot() function computes this operation by taking two tensors and summing their products over the axes you specify.

Syntax

numpy.tensordot(a, b, axes=2)

Parameters

a, b: Input tensors to compute the dot product
axes: Can be an integer N (sum over last N axes of a and first N axes of b) or a tuple of two arrays specifying which axes to sum over

Basic Example

Let's start with a simple example using 2D arrays ?

import numpy as np

# Create two 2D arrays
a = np.array([[1, 2], [3, 4]])
b = np.array([[5, 6], [7, 8]])

print("Array a:")
print(a)
print("\nArray b:")
print(b)

# Tensor dot product with axes=1 (equivalent to matrix multiplication)
result = np.tensordot(a, b, axes=1)
print("\nTensor dot product (axes=1):")
print(result)

Array a:
[[1 2]
 [3 4]]

Array b:
[[5 6]
 [7 8]]

Tensor dot product (axes=1):
[[19 22]
 [43 50]]

Complex Example with 3D Arrays

Here's a more complex example using 3D tensors with specific axis specification ?

import numpy as np

# Creating two numpy 3D arrays
arr1 = np.arange(60.).reshape(3, 4, 5)
arr2 = np.arange(24.).reshape(4, 3, 2)

print("Array1 shape:", arr1.shape)
print("Array2 shape:", arr2.shape)

# Compute tensor dot product with specific axes
# axes=([1,0],[0,1]) means:
# - Sum over axis 1 and 0 of arr1 
# - Sum over axis 0 and 1 of arr2
result = np.tensordot(arr1, arr2, axes=([1,0],[0,1]))

print("\nTensor dot product shape:", result.shape)
print("\nTensor dot product result:")
print(result)

Array1 shape: (3, 4, 5)
Array2 shape: (4, 3, 2)

Tensor dot product shape: (5, 2)

Tensor dot product result:
[[4400. 4730.]
 [4532. 4874.]
 [4664. 5018.]
 [4796. 5162.]
 [4928. 5306.]]

Different Axis Specifications

You can specify axes in different ways ?

import numpy as np

# Create smaller arrays for clarity
a = np.arange(12).reshape(2, 2, 3)
b = np.arange(12).reshape(2, 3, 2)

print("Array a shape:", a.shape)  # (2, 2, 3)
print("Array b shape:", b.shape)  # (2, 3, 2)

# Method 1: Using integer (sum over last N axes of a, first N of b)
result1 = np.tensordot(a, b, axes=1)
print("\nUsing axes=1, result shape:", result1.shape)

# Method 2: Specifying exact axes
result2 = np.tensordot(a, b, axes=([2], [1]))
print("Using axes=([2], [1]), result shape:", result2.shape)

Array a shape: (2, 2, 3)
Array b shape: (2, 3, 2)

Using axes=1, result shape: (2, 2, 2, 2)
Using axes=([2], [1]), result shape: (2, 2, 2, 2)

How It Works

The tensor dot product works by:

• Identifying which axes to sum over in each tensor
• Computing element-wise products across these axes
• Summing the products to reduce dimensionality
• Returning a tensor with the remaining dimensions

Conclusion

NumPy's tensordot() function provides a powerful way to compute generalized dot products between tensors. Use integer axes for simple cases or specify exact axes with tuples for precise control over which dimensions are summed.