Skip to content

← JAX Foundations step 24 of 25

Easy Primitives

Explicit Broadcasting via broadcast_to

Jump to the task

Give a vector a batch shape

Suppose a bias vector contains one value for each of three features. To represent that same vector for four examples, broadcast its shape from (3,) to (4, 3). Every row of the result has the same values.

jnp.broadcast_to(x, shape) makes the output shape explicit. In many calculations you can instead let arithmetic broadcast automatically: adding a (3,) bias to a (4, 3) array already applies it to each row.

Worked example

import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
y = jnp.broadcast_to(x, (4, 3))
# [[1., 2., 3.],
#  [1., 2., 3.],
#  [1., 2., 3.],
#  [1., 2., 3.]]

Shapes align from the right. The existing length-three axis matches the target's last axis. The new leading axis supplies the four rows. More generally, an aligned input dimension must equal the target dimension or be 1. For example, (1,) can broadcast to (2, 3), but (2,) cannot broadcast to (2, 3).

Broadcasting describes values and shapes, not a guaranteed memory layout. Do not assume that a JAX broadcast is a NumPy-style zero-copy view or that it always allocates repeated data. Compilation can change how the result is represented or combine it with later operations.

Problem

Implement explicit_broadcast(x, target_shape) with jnp.broadcast_to.

  • x is a one-dimensional float32 array with shape (W,).
  • target_shape is a float32 array [H, W] of whole-number dimensions. Dimensions are nonnegative and the supplied width matches x.
  • Return an array of shape (H, W). Each of its rows equals x.

For example, [5, 6] with target shape [3, 2] gives [[5, 6], [5, 6], [5, 6]]. Zero rows or zero columns give an empty output with the requested two-dimensional shape.

Shape values and JIT

In this exercise's eager call, convert both shape entries with int and pass the tuple (H, W) to jnp.broadcast_to. This works because the array values are concrete when the function runs.

Wrapping that function directly in jax.jit makes the shape array dynamic. int cannot extract a Python integer from a traced value. For a JIT version, change the interface to accept a Python tuple of integers and mark it static, or fix the target shape in a compiled helper. Making the original JAX array static does not work because it is unhashable.

Broadcasting and repetition

jnp.tile repeats blocks of values. It can turn [1, 2] into [1, 2, 1, 2]; broadcasting cannot expand a length-two axis to four. Choose based on the shape operation you need, not a blanket claim that one function always copies and the other never does.

JAX arrays are immutable. An update such as y.at[0, 0].set(9) returns a new array value; it does not change x or the existing y.

The diagram shows how the vector's trailing axis aligns with the target.

Loading visualization…