jax-ml/jax · error · ValueError

Array.contains: query value must be a scalar, got

Error message

Array.__contains__: query value must be a scalar, got {query.shape=}

What it means

jax.Array.__contains__ requires the query (right operand of `in`) to be a scalar; NumPy allows array-like queries. JAX restricts this to keep semantics simple and to avoid ambiguity about what array-in-array containment means. The message includes the query's shape.

Solutions

Use set-style logic instead: `jnp.isin(arr, query).any()`
Pass a true scalar: `float(query) in arr`
For subset checks use `jnp.isin` directly

Example fix

# before
jnp.array([1., 2.]) in arr

# after
jnp.isin(jnp.array([1., 2.]), arr).all()

Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp

def member(arr, q):
    q = jnp.asarray(q)
    if q.ndim != 0:
        return bool(jnp.isin(q, arr).all())
    return bool(q in arr)

Type guard

def is_scalar_like(q) -> bool:
    return jnp.asarray(q).ndim == 0

Prevention

Don't wrap query values in arrays before `in`
Use jnp.isin for array-vs-array containment

When it happens

Trigger: Calling `query in arr` where `query` is a 0-d-or-higher jnp/np array with ndim != 0, e.g. `jnp.array([1,2]) in jnp_array`.

Common situations: Porting NumPy code that tests array containment; mistakenly wrapping the query value in jnp.array or np.array before the membership test.

Related errors

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

Appendix: source

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

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.
  """
  return reductions.cumprod(self, axis=axis, dtype=dtype, out=out)

View on GitHub (pinned to 1e1c6a8fc0)