Skip to content

Instantly share code, notes, and snippets.

@sorenblank
Created June 20, 2025 19:32
Show Gist options
  • Select an option

  • Save sorenblank/800eb05b8041234f193ff8d8171ffaeb to your computer and use it in GitHub Desktop.

Select an option

Save sorenblank/800eb05b8041234f193ff8d8171ffaeb to your computer and use it in GitHub Desktop.

Multiplying 3D arrays (also known as 3rd-order tensors) is a common operation in fields like scientific computing, machine learning, and data analysis. The term "multiplication" can mean a few different things, so let's break down the most common types.

We'll use Python's NumPy library for the examples, as it's the standard for numerical operations in Python.

First, let's set up two sample 3D arrays. Imagine these represent a "batch" of 2 matrices, where each matrix has 3 rows and 4 columns.

import numpy as np

# Shape: (batch_size, rows, cols) = (2, 3, 4)
a = np.arange(24).reshape(2, 3, 4)
# Shape: (batch_size, rows, cols) = (2, 3, 4)
b = np.arange(24, 48).reshape(2, 3, 4)

print("Array 'a':\n", a)
print("\nArray 'b':\n", b)

1. Element-wise Multiplication

This is the simplest form of multiplication. It multiplies the elements at the exact same position in each array.

Rule: The arrays must have the exact same shape, or be "broadcastable" to the same shape.

How to do it: Use the * operator or the np.multiply() function.

Example

Let's multiply our arrays a and b. The element at position [0, 0, 0] in a (which is 0) will be multiplied by the element at [0, 0, 0] in b (which is 24), and so on for every element.

# Element-wise multiplication
result_elementwise = a * b
# Or equivalently: result_elementwise = np.multiply(a, b)

print("Shape of a:", a.shape)
print("Shape of b:", b.shape)
print("Shape of result:", result_elementwise.shape)
print("\nElement-wise result:\n", result_elementwise)

Multiplying Multiple Arrays: This is straightforward. You can simply chain the operation.

c = np.ones((2, 3, 4)) * 2  # A 3D array of 2s
result_multiple = a * b * c

2. Batched Matrix Multiplication

This is the most common interpretation of "matrix multiplication" for 3D arrays. You treat the 3D array as a "stack" or "batch" of 2D matrices and perform a standard matrix multiplication on each corresponding pair of matrices in the stack.

Rule: For two 3D arrays with shapes ((B, M, K)) and ((B, K, N)):

  • The first dimension ((B), the batch size) must be the same.
  • The inner dimensions ((K)) must match, just like in standard 2D matrix multiplication.
  • The resulting array will have the shape ((B, M, N)).

How to do it: Use the @ operator (recommended) or the np.matmul() function.

Example

We can't multiply a ((2, 3, 4)) by b ((2, 3, 4)) this way because their inner dimensions don't match (4 is not equal to 3). We need to create a new array that is compatible.

# a has shape (2, 3, 4)
# Let's create a compatible array 'c' with shape (2, 4, 5)
c = np.arange(40).reshape(2, 4, 5)

print("Shape of a:", a.shape)
print("Shape of c:", c.shape)

# Batched matrix multiplication
result_matmul = a @ c
# Or equivalently: result_matmul = np.matmul(a, c)

print("\nShape of result:", result_matmul.shape)
print("\nBatched Matmul result:\n", result_matmul)

# To verify, let's look at the first batch item
print("\nFirst matrix from 'a':\n", a[0])
print("\nFirst matrix from 'c':\n", c[0])
print("\nResult of a[0] @ c[0]:\n", a[0] @ c[0])
print("\nFirst matrix from the result:\n", result_matmul[0])

As you can see, result_matmul[0] is the result of the 2D matrix multiplication a[0] @ c[0], and result_matmul[1] is the result of a[1] @ c[1].

Multiplying Multiple Arrays: This works just like chaining matrix multiplication. The shapes must be compatible in sequence.

# a: (2, 3, 4)
# c: (2, 4, 5)
# d: (2, 5, 2)
d = np.arange(20).reshape(2, 5, 2)

# The result will have shape (2, 3, 2)
result_multiple_matmul = a @ c @ d
print("\nShape of a @ c @ d:", result_multiple_matmul.shape)

3. Tensor Dot Product (Tensor Contraction)

This is the most general and powerful form of multiplication. It allows you to specify which axes (dimensions) you want to multiply and sum over. Both element-wise and matrix multiplication are special cases of the tensor dot product.

Rule: The lengths of the axes you specify for contraction must be equal.

How to do it: Use the np.tensordot() function.

Example

Let's replicate our batched matrix multiplication using tensordot. We want to multiply a (shape ((2, 3, 4))) and c (shape ((2, 4, 5))).

For a batched matmul, we need to:

  1. Match the batch dimensions (axis 0 of a with axis 0 of c).
  2. Multiply and sum over the inner matrix dimensions (axis 2 of a with axis 1 of c).
# a: (2, 3, 4)
# c: (2, 4, 5)
# axes=([0, 2], [0, 1]) means:
# - contract axis 0 of 'a' with axis 0 of 'c'
# - contract axis 2 of 'a' with axis 1 of 'c'
result_tensordot = np.tensordot(a, c, axes=([0, 2], [0, 1]))

print("Shape of a:", a.shape)
print("Shape of c:", c.shape)
print("Shape of tensordot result:", result_tensordot.shape)

# Note: The result is not in the same order as matmul.
# The resulting axes are (remaining_axes_of_a, remaining_axes_of_c)
# which is (axis 1 of a, axis 2 of c) -> shape (3, 5)
# This is not what we want for a batched matmul.
# np.matmul is much better for that specific task.

A more common use for tensordot is for different kinds of contractions. For example, let's contract the last axis of a ((2, 3, 4)) with the first axis of a new array e ((4, 6)).

e = np.arange(24).reshape(4, 6)

# Contract axis 2 of 'a' (length 4) with axis 0 of 'e' (length 4)
# axes=([2], [0])
result_td_2 = np.tensordot(a, e, axes=([2], [0]))

print("Shape of a:", a.shape)
print("Shape of e:", e.shape)
print("Shape of result:", result_td_2.shape)
# The result shape is the concatenation of the remaining axes: (2, 3, 6)

Summary: Which One Should You Use?

Multiplication Type When to Use NumPy Operator/Function
Element-wise You want to multiply corresponding elements. Useful for scaling, applying masks, or simple arithmetic. * or np.multiply()
Batched Matrix You have a stack/batch of matrices and want to perform matrix multiplication on each pair. Very common in deep learning. @ or np.matmul()
Tensor Dot Product You need a more general multiplication, specifying exactly which dimensions to sum over. The most flexible and powerful option. np.tensordot() or np.einsum() (even more advanced)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment