jax-ml/jax · error · ValueError

np.delete(arr, obj): for boolean indices, obj must be one-di

Error message

np.delete(arr, obj): for boolean indices, obj must be one-dimensional with length matching specified axis.

What it means

Raised by jnp.delete when the boolean mask obj does not have exactly shape (a.shape[axis],). JAX implements delete by building a boolean keep-mask over the given axis, so a boolean obj must be 1-D and match the axis length exactly, unlike integer indices which can be arbitrary.

Source

Thrown at jax/_src/numpy/lax_numpy.py:7724

        0,
        a.shape[axis],
    )
    obj_array = sort(obj_array)
    obj_array -= arange(len(obj_array), dtype=obj_array.dtype)
    i = arange(a.shape[axis] - obj_array.size, dtype=obj_array.dtype)
    i += (i[None, :] >= obj_array[:, None]).sum(0, dtype=i.dtype)
    return a[(slice(None),) * axis + (i,)]

  # Case 3b: non-unique indices: must be static.
  obj_array = core.concrete_or_error(np.asarray, obj, "'obj' array argument of jnp.delete()")
  if issubdtype(obj_array.dtype, np.integer):
    # TODO(jakevdp): in theory this could be done dynamically if obj has no duplicates,
    # but this would require the complement of lax.gather.
    mask = np.ones(a.shape[axis], dtype=bool)
    mask[obj_array] = False
  elif obj_array.dtype == bool:
    if obj_array.shape != (a.shape[axis],):
      raise ValueError("np.delete(arr, obj): for boolean indices, obj must be one-dimensional "
                       "with length matching specified axis.")
    mask = ~obj_array
  else:
    raise ValueError(f"np.delete(arr, obj): got obj.dtype={obj_array.dtype}; must be integer or bool.")
  return a[tuple(slice(None) for i in range(axis)) + (mask,)]


@export
def insert(arr: ArrayLike, obj: ArrayLike | slice, values: ArrayLike,
           axis: int | None = None) -> Array:
  """Insert entries into an array at specified indices.

  JAX implementation of :func:`numpy.insert`.

  Args:
    arr: array object into which values will be inserted.
    obj: slice or array of indices specifying insertion locations.
    values: array of values to be inserted.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/flatten the mask to 1-D with length equal to a.shape[axis] (e.g. mask.ravel() or mask[:, 0])
2. Pass the correct axis argument matching the mask's length
3. Convert the mask to integer indices: jnp.delete(a, jnp.where(mask)[0], axis=k)

Example fix

// before
jnp.delete(a, mask, axis=1)  # mask has shape (a.shape[0], 1)
// after
jnp.delete(a, mask.ravel(), axis=0)  # or mask[:, 0] if mask is (n,1) targeting columns

mask = jnp.asarray(mask)
assert mask.ndim == 1 and mask.shape[0] == a.shape[axis], 'mask must match axis length'

def is_valid_delete_mask(a, mask, axis=0):
    m = jnp.asarray(mask)
    return m.dtype == jnp.bool_ and m.ndim == 1 and m.shape[0] == a.shape[axis]

When it happens

Trigger: Calling jnp.delete(a, mask, axis=k) where mask is boolean with shape != (a.shape[k],), e.g. a 2-D boolean array, a mask sized for a different axis, or the default axis=0 while the mask was built for another axis.

Common situations: Passing a column mask (shape (n,1)) instead of a flat mask; forgetting to specify axis when deleting rows vs columns; porting NumPy code where a differently-sized boolean was silently tolerated via integer conversion.

Related errors

• boolean index did not match shape of indexed array in index
• 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
• Incorrect output shape for return value #{i}: Expected: {out

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