jax-ml/jax · error · IndexError

boolean index did not match shape of indexed array in index

Error message

boolean index did not match shape of indexed array in index {position}: got {idx_shape}, expected {expected_shape}

What it means

expand_bool_indices verifies each boolean mask's shape equals the shape of the axes it indexes (each dim must match or be 0). Like NumPy, a boolean mask must have the same length as the dimensions it covers; otherwise this IndexError is raised with got/expected shapes.

Source

Thrown at jax/_src/numpy/indexing.py:286

    expanded_indices: list[ParsedIndex] = []

    for position, idx in enumerate(self.indices):
      if idx.typ != IndexType.BOOLEAN:
        expanded_indices.append(idx)
        continue
      if not core.is_concrete(idx.index):
        # TODO(mattjj): improve this error by tracking _why_ the indices are not concrete
        raise errors.NonConcreteBooleanIndexError(core.typeof(idx.index))
      assert isinstance(idx.index, (bool, np.ndarray, Array, list))
      if np.ndim(idx.index) == 0:  # pyrefly: ignore[bad-argument-type]
        # Scalar booleans
        assert idx.consumed_axes == ()
        expanded_indices.append(ParsedIndex(index=bool(idx.index), typ=idx.typ, consumed_axes=()))
        continue
      idx_shape = np.shape(idx.index)  # pyrefly: ignore[no-matching-overload]
      expected_shape = [self.shape[i] for i in idx.consumed_axes]
      if not all(s1 in (0, s2) for s1, s2 in zip(idx_shape, expected_shape)):
        raise IndexError("boolean index did not match shape of indexed array in index"
                        f" {position}: got {idx_shape}, expected {expected_shape}")
      expanded_indices_raw = np.where(np.asarray(idx.index))
      expanded_indices.extend(ParsedIndex(index=i, typ=IndexType.ARRAY, consumed_axes=(axis,))
                              for i, axis in safe_zip(expanded_indices_raw, idx.consumed_axes))
    return NDIndexer(shape=self.shape, indices=expanded_indices)

  def expand_scalar_bool_indices(self, sharding_spec: Any = None) -> tuple[NDIndexer, Any]:
    new_shape = list(self.shape)
    new_sharding_spec = list((None for _ in self.shape) if sharding_spec is None else sharding_spec.partitions)
    new_indices = list(self.indices)
    current_dim = 0
    for i, idx in enumerate(self.indices):
      if idx.typ == IndexType.BOOLEAN and np.ndim(idx.index) == 0:  # pyrefly: ignore[bad-argument-type]
        new_shape.insert(i, 1)
        new_sharding_spec.insert(i, None)
        new_indices[i] = ParsedIndex(
          np.arange(int(idx.index)), typ=IndexType.ARRAY, consumed_axes=(current_dim,))  # pyrefly: ignore[bad-argument-type]
        current_dim += 1

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rebuild the mask from the current array: mask = x[:, 0] > 0
2. Reshape/reduce mask to the indexed axes: mask = mask[: x.shape[0]] only if logically correct
3. Verify shapes: assert mask.shape == x.shape before masking

Example fix

// before
mask = other > 0
y = x[mask]
// after
mask = x.sum(axis=1) > 0
y = x[mask]

mask = jnp.asarray(mask, dtype=bool)
assert all(m in (0, s) for m, s in zip(mask.shape, expected_shape))

def mask_matches(mask, arr) -> bool:
    return jnp.asarray(mask).shape == arr.shape

When it happens

Trigger: x = jnp.zeros((3,4)); mask = jnp.array([True,False]); x[mask] — mask length 2 vs axis size 3; or a mask covering one axis but shaped for another.

Common situations: Mask built from a different array or stale shape (e.g. train/test split change), mask computed on flattened data applied to 2-D array, or transposed arrays.

Related errors

• Too many indices: array is {len(shape)}-dimensional, but {to
• np.delete(arr, obj): for boolean indices, obj must be one-di
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac

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