jax-ml/jax · error · ValueError

fill_value shape {fv_aval.shape} cannot be broadcast to batc

Error message

fill_value shape {fv_aval.shape} cannot be broadcast to batch shape {batch_shape} without expanding it.

What it means

Raised by the Nonzero HiJAX primitive when a fill_value broadcasts against the batch shape only by expanding it (adding leading dims), i.e. broadcasted != batch_shape. The primitive disallows expanding fills: the broadcast result must equal the batch shape exactly, so shapes like (1, 4) against batch (4,) are rejected even though numpy would allow them.

Source

Thrown at jax/_src/numpy/hijax.py:224

    if len(axes) != len(set(axes)):
      raise ValueError(f"duplicate axes are not allowed: {axes=}")
    if fill_value_avals and len(fill_value_avals) != len(axes):
      raise ValueError(f"Expected {len(axes)} fill values, got {len(fill_value_avals)}")
    if any(fv.dtype != out_dtype for fv in fill_value_avals):
      raise ValueError(f"Expected fill values to have dtype {out_dtype}, got {fill_value_avals}")
    batch_shape = tuple(
        s for i, s in enumerate(a_aval.shape) if i not in axes
    )
    for fv_aval in fill_value_avals:
      try:
        broadcasted = lax.broadcast_shapes(fv_aval.shape, batch_shape)
      except ValueError as e:
        raise ValueError(
            f"fill_value shape {fv_aval.shape} is not broadcast-compatible with "
            f"batch shape {batch_shape}"
        ) from e
      if broadcasted != batch_shape:
        raise ValueError(
            f"fill_value shape {fv_aval.shape} cannot be broadcast to "
            f"batch shape {batch_shape} without expanding it."
        )
    self.in_avals = (a_aval, *fill_value_avals)

    # Evaluate shape to set out_aval
    self.out_aval = tree_util.tree_map(core.typeof, api.eval_shape(
        functools.partial(_nonzero_impl, size=size, axes=axes, out_dtype=out_dtype),
        a_aval, *fill_value_avals))

    self.params = dict(
        size=size,
        axes=axes,
        out_dtype=out_dtype,
    )
    super().__init__()

  def expand(self, a: ArrayLike, *fill_value: ArrayLike) -> tuple[Array, ...]:  # pyrefly: ignore[bad-override]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Squeeze the fill value to remove extra leading dims: fv = fv.reshape(batch_shape)
Prefer scalar fill values
If expansion is genuinely needed, broadcast the fill value explicitly to batch_shape before constructing the primitive

Example fix

# before
fv = jnp.zeros((1, 8))  # batch shape is (8,)
# after
fv = jnp.zeros((8,))  # or jnp.zeros(())

Defensive patterns

Strategy: validation

Validate before calling

from jax._src.lax import lax
b = lax.broadcast_shapes(fv.shape, batch_shape)
assert b == tuple(batch_shape), (fv.shape, batch_shape)

Type guard

def fill_exact_broadcast(fv, batch_shape) -> bool:
    from jax._src.lax import lax
    return lax.broadcast_shapes(fv.shape, tuple(batch_shape)) == tuple(batch_shape)

Prevention

Reshape fills to the batch shape: fv = fv.reshape(batch_shape)
Avoid stray leading size-1 dims; prefer scalars
This API is stricter than numpy broadcasting — expanding fills is banned

When it happens

Trigger: A fill value of shape (1, B) with batch shape (B,) — broadcasting succeeds but yields (1, B), so the primitive raises. Scalar shapes are unaffected.

Common situations: Squeezing/reshaping fill values so they carry a stray leading 1 dim; migrating code that relied on implicit dimension expansion.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/e3ee5124b074d8dd. Report an issue: GitHub.