How to beginner · 3 min read

How to use numpy broadcasting

Q: How to use numpy broadcasting

Use numpy broadcasting to perform element-wise operations on arrays of different shapes by automatically expanding their dimensions. In PyTorch, broadcasting works similarly, allowing tensors with compatible shapes to be combined without explicit reshaping.

Quick answer

Use numpy broadcasting to perform element-wise operations on arrays of different shapes by automatically expanding their dimensions. In PyTorch, broadcasting works similarly, allowing tensors with compatible shapes to be combined without explicit reshaping.

PREREQUISITES

Python 3.8+
pip install torch>=2.0.0
pip install numpy>=1.24

Setup

Install PyTorch and NumPy if you haven't already. Use the following commands to install them via pip:

bash

pip install torch numpy

Step by step

This example demonstrates numpy broadcasting principles applied to PyTorch tensors. Broadcasting lets you add tensors of different shapes without explicit loops or reshaping.

python

import torch

# Create a tensor of shape (3, 1)
a = torch.tensor([[1], [2], [3]])  # shape: (3,1)

# Create a tensor of shape (1, 4)
b = torch.tensor([10, 20, 30, 40])  # shape: (1,4)

# Add tensors using broadcasting
c = a + b

print('Tensor a shape:', a.shape)
print('Tensor b shape:', b.shape)
print('Result c shape:', c.shape)
print('Result c values:\n', c)

output

Tensor a shape: torch.Size([3, 1])
Tensor b shape: torch.Size([1, 4])
Result c shape: torch.Size([3, 4])
Result c values:
 tensor([[11, 21, 31, 41],
        [12, 22, 32, 42],
        [13, 23, 33, 43]])

Common variations

You can use broadcasting with other operations like multiplication, subtraction, and division. It also works with higher-dimensional tensors as long as trailing dimensions are compatible.

Example: multiplying a (2,3,1) tensor with a (1,3,4) tensor broadcasts to (2,3,4).

python

x = torch.ones((2, 3, 1))  # shape (2,3,1)
y = torch.arange(12).reshape(1, 3, 4)  # shape (1,3,4)

z = x * y

print('x shape:', x.shape)
print('y shape:', y.shape)
print('z shape:', z.shape)
print('z values:\n', z)

output

x shape: torch.Size([2, 3, 1])
y shape: torch.Size([1, 3, 4])
z shape: torch.Size([2, 3, 4])
z values:
 tensor([[[ 0.,  1.,  2.,  3.],
          [ 4.,  5.,  6.,  7.],
          [ 8.,  9., 10., 11.]],

         [[ 0.,  1.,  2.,  3.],
          [ 4.,  5.,  6.,  7.],
          [ 8.,  9., 10., 11.]]])

Troubleshooting

If you get a RuntimeError: The size of tensor a (X) must match the size of tensor b (Y) at non-singleton dimension Z, it means the tensors are not broadcast-compatible. Check that dimensions either match or one of them is 1.

Use tensor.shape to debug shapes before operations.

✅

Key Takeaways

Broadcasting lets you perform element-wise operations on tensors with different but compatible shapes without explicit reshaping.
In PyTorch, broadcasting follows the same rules as NumPy, expanding dimensions of size 1 to match the other tensor.
Always verify tensor shapes before operations to avoid runtime shape mismatch errors.

Verified 2026-04

Verify ↗