jax-ml/jax · error · ValueError

fill_value shape cannot be broadcast to batch shape without…

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.

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.

Appendix: 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)