Multiply Each Last-Axis Slice in NumPy with Broadcasting: Add a Trailing Axis via y[..., np.newaxis]

Learn how to multiply last-axis slices in NumPy by broadcasting y with y[..., np.newaxis], avoiding loops and preallocation. Clear axis alignment tips included.

When working with NumPy broadcasting, a frequent pattern is multiplying each slice along the last axis of one array by another array whose shape broadcasts with all but that last axis. The straightforward approach uses a loop and preallocation, but there is a more idiomatic way to express the same operation with proper axis alignment.

Problem setup

Assume two arrays, where x.shape[:-1] and y.shape are broadcastable. The goal is to produce z so that each slice x[..., k] is multiplied by y and written into z[..., k]. A direct implementation might look like this:

import numpy as np
x = ...  # some ndarray
y = ...  # some ndarray with shape broadcastable to x.shape[:-1]
z = np.empty(np.broadcast_shapes(x.shape[:-1], y.shape), x.dtype)
for k in range(x.shape[-1]):
    z[..., k] = x[..., k] * y

What is really going on

This operation is essentially broadcasting y across the last dimension of x. The loop is there only to work around the fact that y lacks that extra dimension. If y had a length-1 dimension appended at the end, NumPy would broadcast it over the final axis of x automatically. The placement of this new axis is crucial: putting it at the end aligns with the last axis of x; putting it at the beginning does not. Mixing up these positions can lead to “operands could not be broadcast together” errors.

Concise solution

Use np.newaxis together with ... to add a singleton dimension after the last axis of y and rely on broadcasting for the rest. This eliminates the loop and does not require preallocating the output.

import numpy as np
x = ...  # same x as above
y = ...  # same y as above
z = x * y[..., np.newaxis]

Note that adding the axis at the end with y[..., np.newaxis] is not the same as placing it at the start. The position matters for broadcasting.

Why this matters

Expressing the intent via broadcasting makes the code shorter and clearer. It avoids manual allocation of the output array and removes the explicit Python-level loop, while preserving the exact computation. The key is to add the new axis after the last axis of y so that it broadcasts over the final dimension of x.

Takeaways

When you need to multiply each last-axis slice of one array by another array whose shape is broadcastable with all preceding axes, add a trailing singleton dimension to the second array: y[..., np.newaxis]. This aligns shapes correctly for broadcasting, removes boilerplate, and keeps the operation idiomatic to NumPy.