{"record":{"id":"e3ee5124b074d8dd","repo":"jax-ml/jax","slug":"fill-value-shape-fv-aval-shape-cannot-be-broadca","errorCode":null,"errorMessage":"fill_value shape {fv_aval.shape} cannot be broadcast to batch shape {batch_shape} without expanding it.","messagePattern":"fill_value shape (.+?) cannot be broadcast to batch shape (.+?) without expanding it\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/hijax.py","lineNumber":224,"sourceCode":"    if len(axes) != len(set(axes)):\n      raise ValueError(f\"duplicate axes are not allowed: {axes=}\")\n    if fill_value_avals and len(fill_value_avals) != len(axes):\n      raise ValueError(f\"Expected {len(axes)} fill values, got {len(fill_value_avals)}\")\n    if any(fv.dtype != out_dtype for fv in fill_value_avals):\n      raise ValueError(f\"Expected fill values to have dtype {out_dtype}, got {fill_value_avals}\")\n    batch_shape = tuple(\n        s for i, s in enumerate(a_aval.shape) if i not in axes\n    )\n    for fv_aval in fill_value_avals:\n      try:\n        broadcasted = lax.broadcast_shapes(fv_aval.shape, batch_shape)\n      except ValueError as e:\n        raise ValueError(\n            f\"fill_value shape {fv_aval.shape} is not broadcast-compatible with \"\n            f\"batch shape {batch_shape}\"\n        ) from e\n      if broadcasted != batch_shape:\n        raise ValueError(\n            f\"fill_value shape {fv_aval.shape} cannot be broadcast to \"\n            f\"batch shape {batch_shape} without expanding it.\"\n        )\n    self.in_avals = (a_aval, *fill_value_avals)\n\n    # Evaluate shape to set out_aval\n    self.out_aval = tree_util.tree_map(core.typeof, api.eval_shape(\n        functools.partial(_nonzero_impl, size=size, axes=axes, out_dtype=out_dtype),\n        a_aval, *fill_value_avals))\n\n    self.params = dict(\n        size=size,\n        axes=axes,\n        out_dtype=out_dtype,\n    )\n    super().__init__()\n\n  def expand(self, a: ArrayLike, *fill_value: ArrayLike) -> tuple[Array, ...]:  # pyrefly: ignore[bad-override]","sourceCodeStart":206,"sourceCodeEnd":242,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/hijax.py#L206-L242","documentation":"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.","triggerScenarios":"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.","commonSituations":"Squeezing/reshaping fill values so they carry a stray leading 1 dim; migrating code that relied on implicit dimension expansion.","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"],"exampleFix":"# before\nfv = jnp.zeros((1, 8))  # batch shape is (8,)\n# after\nfv = jnp.zeros((8,))  # or jnp.zeros(())","handlingStrategy":"validation","validationCode":"from jax._src.lax import lax\nb = lax.broadcast_shapes(fv.shape, batch_shape)\nassert b == tuple(batch_shape), (fv.shape, batch_shape)","typeGuard":"def fill_exact_broadcast(fv, batch_shape) -> bool:\n    from jax._src.lax import lax\n    return lax.broadcast_shapes(fv.shape, tuple(batch_shape)) == tuple(batch_shape)","tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","nonzero","fill-value","broadcasting"],"backgroundTag":"broadcast-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}