{"record":{"id":"ee7ab752c604c56d","repo":"jax-ml/jax","slug":"boolean-index-did-not-match-shape-of-indexed-array","errorCode":null,"errorMessage":"boolean index did not match shape of indexed array in index {position}: got {idx_shape}, expected {expected_shape}","messagePattern":"boolean index did not match shape of indexed array in index (.+?): got (.+?), expected (.+?)","errorType":"validation","errorClass":"IndexError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":286,"sourceCode":"    expanded_indices: list[ParsedIndex] = []\n\n    for position, idx in enumerate(self.indices):\n      if idx.typ != IndexType.BOOLEAN:\n        expanded_indices.append(idx)\n        continue\n      if not core.is_concrete(idx.index):\n        # TODO(mattjj): improve this error by tracking _why_ the indices are not concrete\n        raise errors.NonConcreteBooleanIndexError(core.typeof(idx.index))\n      assert isinstance(idx.index, (bool, np.ndarray, Array, list))\n      if np.ndim(idx.index) == 0:  # pyrefly: ignore[bad-argument-type]\n        # Scalar booleans\n        assert idx.consumed_axes == ()\n        expanded_indices.append(ParsedIndex(index=bool(idx.index), typ=idx.typ, consumed_axes=()))\n        continue\n      idx_shape = np.shape(idx.index)  # pyrefly: ignore[no-matching-overload]\n      expected_shape = [self.shape[i] for i in idx.consumed_axes]\n      if not all(s1 in (0, s2) for s1, s2 in zip(idx_shape, expected_shape)):\n        raise IndexError(\"boolean index did not match shape of indexed array in index\"\n                        f\" {position}: got {idx_shape}, expected {expected_shape}\")\n      expanded_indices_raw = np.where(np.asarray(idx.index))\n      expanded_indices.extend(ParsedIndex(index=i, typ=IndexType.ARRAY, consumed_axes=(axis,))\n                              for i, axis in safe_zip(expanded_indices_raw, idx.consumed_axes))\n    return NDIndexer(shape=self.shape, indices=expanded_indices)\n\n  def expand_scalar_bool_indices(self, sharding_spec: Any = None) -> tuple[NDIndexer, Any]:\n    new_shape = list(self.shape)\n    new_sharding_spec = list((None for _ in self.shape) if sharding_spec is None else sharding_spec.partitions)\n    new_indices = list(self.indices)\n    current_dim = 0\n    for i, idx in enumerate(self.indices):\n      if idx.typ == IndexType.BOOLEAN and np.ndim(idx.index) == 0:  # pyrefly: ignore[bad-argument-type]\n        new_shape.insert(i, 1)\n        new_sharding_spec.insert(i, None)\n        new_indices[i] = ParsedIndex(\n          np.arange(int(idx.index)), typ=IndexType.ARRAY, consumed_axes=(current_dim,))  # pyrefly: ignore[bad-argument-type]\n        current_dim += 1","sourceCodeStart":268,"sourceCodeEnd":304,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L268-L304","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Rebuild the mask from the current array: mask = x[:, 0] > 0","Reshape/reduce mask to the indexed axes: mask = mask[: x.shape[0]] only if logically correct","Verify shapes: assert mask.shape == x.shape before masking"],"exampleFix":"// before\nmask = other > 0\ny = x[mask]\n// after\nmask = x.sum(axis=1) > 0\ny = x[mask]","handlingStrategy":"validation","validationCode":"mask = jnp.asarray(mask, dtype=bool)\nassert all(m in (0, s) for m, s in zip(mask.shape, expected_shape))","typeGuard":"def mask_matches(mask, arr) -> bool:\n    return jnp.asarray(mask).shape == arr.shape","tryCatchPattern":null,"preventionTips":["Build masks from the array being indexed","Pin/verify data shapes at pipeline boundaries"],"tags":["jax","boolean-mask","shape-mismatch","indexing"],"backgroundTag":"boolean-mask-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}