jax-ml/jax · error · ValueError

fill_value shape {fv_aval.shape} is not broadcast-compatible

Error message

fill_value shape {fv_aval.shape} is not broadcast-compatible with batch shape {batch_shape}

What it means

Raised by the Nonzero HiJAX primitive when a fill_value's shape cannot be broadcast together with the batch shape (the input shape with the reduced axes removed). lax.broadcast_shapes raised a ValueError, which is re-raased with this clearer message. Fill values must be scalars or broadcast-compatible with the batch dims.

Source

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

        raise ValueError(f"size must be a positive integer; got {size=}")
    if not dtypes.issubdtype(out_dtype, np.integer):
      raise ValueError(f"out_dtype must be integer typed; got {out_dtype=}")
    if not all(0 <= ax < a_aval.ndim for ax in axes):
      raise ValueError(f"axes out of range for array with {a_aval.ndim} dimensions:  {axes=}")
    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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use scalar fill values (shape ()), which always broadcast
  2. Reshape fill values to be broadcast-compatible with batch_shape, or to exactly batch_shape

Example fix

# before
fv = jnp.zeros((3,))  # batch shape is (4,)
# after
fv = jnp.zeros(())  # scalar broadcasts to any batch shape
Defensive patterns

Strategy: validation

Validate before calling

try:
    from jax._src.lax import lax
    lax.broadcast_shapes(fv.shape, batch_shape)
except ValueError:
    raise SystemExit('fill_value incompatible with batch shape')

Type guard

def fill_broadcastable(fv_shape, batch_shape) -> bool:
    try:
        from jax._src.lax import lax
        return lax.broadcast_shapes(tuple(fv_shape), tuple(batch_shape)) is not None
    except ValueError:
        return False

Prevention

When it happens

Trigger: Supplying a fill_value with shape (3,) when the batch shape is (4,) — the shapes conflict and cannot broadcast; or non-scalar fill values with incompatible trailing dims in direct primitive construction.

Common situations: Passing per-batch fill values computed from an array shaped against an older layout; broadcasting rules violated because dims of size 1 were expected.

Related errors

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