jax-ml/jax · error · TypeError

random_unwrap takes key array operand, got {keys.dtype=}

Error message

random_unwrap takes key array operand, got {keys.dtype=}

What it means

jax._src.random.prng.random_unwrap is the internal inverse of random_wrap: it strips the typed-key wrapper and returns raw uint32 key data. It requires its operand's dtype to be a prng_key subtype; passing a plain uint32 array (already unwrapped) or any non-key array raises this TypeError.

Source

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

def random_wrap_impl(base_arr, *, impl):
  return PRNGKeyArray(impl, base_arr)

def random_wrap_lowering(ctx, base_arr, *, impl):
  return [base_arr]

def random_wrap_batch_rule(batched_args, batch_dims, *, impl):
  x, = batched_args
  d, = batch_dims
  x = batching.bdim_at_front(x, d, 1)
  return random_wrap(x, impl=impl), 0

mlir.register_lowering(random_wrap_p, random_wrap_lowering)
batching.primitive_batchers[random_wrap_p] = random_wrap_batch_rule


def random_unwrap(keys):
  if not dtypes.issubdtype(keys.dtype, dtypes.prng_key):
    raise TypeError(f'random_unwrap takes key array operand, got {keys.dtype=}')
  return random_unwrap_p.bind(keys)

random_unwrap_p = core.Primitive('random_unwrap')
ad.defjvp_zero(random_unwrap_p)
batching.defvectorized(random_unwrap_p)

@random_unwrap_p.def_abstract_eval
def random_unwrap_abstract_eval(keys_aval):
  return core.physical_aval(keys_aval)

@random_unwrap_p.def_impl
def random_unwrap_impl(keys):
  return keys._base_array

def random_unwrap_lowering(ctx, keys):
  return [keys]

mlir.register_lowering(random_unwrap_p, random_unwrap_lowering)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Only unwrap typed keys: guard with jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key) before calling
2. If you already have raw data, skip the unwrap
3. Use the public jax.random.key_data instead of internal random_unwrap

Example fix

// before
raw = jax.random.key_data(already_raw_uint32_array)

// after
import jax, jax.numpy as jnp
raw = (jax.random.key_data(x)
       if jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key)
       else jnp.asarray(x, dtype=jnp.uint32))

import jax
if jax.dtypes.issubdtype(keys.dtype, jax.dtypes.prng_key):
    raw = jax.random.key_data(keys)
else:
    raw = keys  # already unwrapped

import jax
def is_typed_key(x) -> bool:
    return jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key)

When it happens

Trigger: Calling jax.random.key_data on an array that is already raw uint32; feeding random_unwrap the output of another random_unwrap; passing a regular JAX array where a typed key is expected.

Common situations: Double-unwrap bugs in serialization round-trips; helper functions that accept 'key or key_data' and unconditionally unwrap; mixing the public jax.random.key_data API with internal prng functions.

Related errors

• JAX encountered invalid PRNG key data: expected key_data.dty
• requires 8-, 16-, 32- or 64-bit field width.
• PRNG key seed must be an integer; got {seed!r}
• threefry_2x32 requires uint32 arguments, got {}
• requires 8-, 16-, 32- or 64-bit field width.

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