jax-ml/jax · error · TypeError

key array cannot be converted to boolean.

Error message

key array cannot be converted to boolean.

What it means

PRNGKeyArray deliberately disables truthiness (__bool__ raises) because there is no meaningful boolean interpretation of a random key. Any use of a key in a boolean context — if key:, key or other, not key, assert key — triggers this TypeError. This is a design decision to prevent silently treating keys as truthy/falsy objects.

Source

Thrown at jax/_src/random/prng.py:311

      raise TypeError('len() of unsized object')
    return len(self._base_array)

  def __iter__(self) -> Iterator[PRNGKeyArray]:
    if self._is_scalar():
      raise TypeError('iteration over a 0-d key array')
    # TODO(frostig): we may want to avoid iteration by slicing because
    # a very common use of iteration is `k1, k2 = split(key)`, and
    # slicing/indexing may be trickier to track for linearity checking
    # purposes. Maybe we can:
    # * introduce an unpack primitive+traceable (also allow direct use)
    # * unpack upfront into shape[0] many keyarray slices
    # * return iter over these unpacked slices
    # Whatever we do, we'll want to do it by overriding
    # ShapedArray._iter when the element type is KeyTy...
    return (PRNGKeyArray(self._impl, k) for k in iter(self._base_array))

  def __bool__(self):
    raise TypeError("key array cannot be converted to boolean.")

  def __repr__(self):
    return (f'Array({self.shape}, dtype={self.dtype.name}) overlaying:\n'
            f'{self._base_array}')

  def pprint(self):
    pp_keys = pp.text('shape = ') + pp.text(str(self.shape))
    pp_impl = pp.text('impl = ') + self._impl.pprint()
    return str(pp.group(
      pp.text('PRNGKeyArray:') +
      pp.nest(2, pp.brk() + pp_keys + pp.brk() + pp_impl)))

  def copy(self):
    out = self.__class__(self._impl, self._base_array.copy())
    out._consumed = self._consumed  # TODO(jakevdp): is this correct?
    return out

  __hash__ = None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace truthiness with an explicit None check: if key is None: key = jax.random.key(0)
2. Use isinstance(key, jax.Array) and jax.dtypes.issubdtype(key.dtype, jax.dtypes.prng_key) to validate
3. Never use keys directly in boolean expressions

Example fix

// before
key = key or jax.random.key(0)  # TypeError

// after
key = key if key is not None else jax.random.key(0)

if key is None:
    key = jax.random.key(0)

import jax, jax.numpy as jnp
def is_typed_key(x) -> bool:
    return isinstance(x, jax.Array) and jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key)

When it happens

Trigger: if key: ...; bool(key); using a key as a condition in 'key or default' expressions; assert key in tests; filter/map predicates that receive keys.

Common situations: Optional-argument patterns like 'key = key or jax.random.key(0)' ; truthiness checks in generic validation helpers; notebook-style asserts on objects.

Related errors

• bool() not supported for instances of type '{0}' (did you me
• lax.bitcast_convert_type does not support bool or complex va
• PRNG keys must be loaded from SMEM. Did you set the memory s
• Seed key_data must be 1D.
• Leading dimension of seed key_data must be 1.

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