jax-ml/jax · error · ValueError

Array.contains: search array must be one-dimensional, go

Error message

Array.__contains__: search array must be one-dimensional, got arr.shape={self.shape}.

What it means

jax.Array.__contains__ requires the array being searched to be one-dimensional, unlike NumPy which flattens. This is a deliberate JAX design choice to avoid NumPy's surprising `in` behavior on multi-dimensional arrays. The error reports the actual shape of the array.

Source

Thrown at jax/_src/numpy/array_methods.py:217

  Refer to :func:`jax.numpy.conjugate` for the full documentation.
  """
  return ufuncs.conjugate(self)

def _contains(self: Array, other: ArrayLike) -> Array:
  """Implements __contains__ for JAX arrays.

  This is used by the Python ``in`` operator.
  """
  # Note: we deliberately depart from NumPy's behavior here, which includes
  # some oddities (https://github.com/numpy/numpy/issues/21933). Namely, we
  # require `self` to be a 1D array, and require `other` to be a scalar.'

  # Explicitly check for string and None types, as these were common bugs.
  if other is None or isinstance(other, str):
    raise TypeError(f"Array.__contains__: unsupported operand type {type(other)}.")
  query = util.ensure_arraylike('Array.__contains__', other)
  if self.ndim != 1:
    raise ValueError("Array.__contains__: search array must be one-dimensional,"
                     f" got arr.shape={self.shape}.")
  if query.ndim != 0:
    raise ValueError("Array.__contains__: query value must be a scalar,"
                     f" got {query.shape=}")
  return reductions.any(self == query)

def _copy(self: Array) -> Array:
  """Return a copy of the array.

  Refer to :func:`jax.numpy.copy` for the full documentation.
  """
  return lax_numpy.copy(self)

def _cumprod(self: Array, axis: int | None = None,
             dtype: DTypeLike | None = None, out: None = None) -> Array:
  """Return the cumulative product of the array.

  Refer to :func:`jax.numpy.cumprod` for the full documentation.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Flatten first: `scalar in arr.flatten()` or `scalar in arr.ravel()`
Index/select a 1D slice: `scalar in arr[i]`
Use an explicit reduction: `bool(jnp.any(arr == scalar))`

Example fix

# before
3.0 in matrix  # matrix.shape == (3, 4)

# after
3.0 in matrix.flatten()

Defensive patterns

Strategy: validation

Validate before calling

def contains_1d(arr, x):
    assert arr.ndim == 1, f'expected 1D, got {arr.shape}'
    return bool(x in arr)

Type guard

def is_1d(a) -> bool:
    return getattr(a, 'ndim', -1) == 1

Prevention

Prefer jnp.any(arr == x) which works for any ndim
Flatten before membership tests when porting NumPy code

When it happens

Trigger: Executing `scalar in arr` where `arr` is a jax.Array with ndim != 1 (e.g. shape (3, 4) or ()).

Common situations: Reusing NumPy code that used `in` on matrices; forgetting to select a row/column (`arr[i] in matrix` vs `x in matrix[i]`); applying membership tests to batched tensors.

Related errors

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