{"record":{"id":"0151dffcb19345bb","repo":"jax-ml/jax","slug":"random-unwrap-takes-key-array-operand-got-keys-d","errorCode":null,"errorMessage":"random_unwrap takes key array operand, got {keys.dtype=}","messagePattern":"random_unwrap takes key array operand, got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/prng.py","lineNumber":777,"sourceCode":"def random_wrap_impl(base_arr, *, impl):\n  return PRNGKeyArray(impl, base_arr)\n\ndef random_wrap_lowering(ctx, base_arr, *, impl):\n  return [base_arr]\n\ndef random_wrap_batch_rule(batched_args, batch_dims, *, impl):\n  x, = batched_args\n  d, = batch_dims\n  x = batching.bdim_at_front(x, d, 1)\n  return random_wrap(x, impl=impl), 0\n\nmlir.register_lowering(random_wrap_p, random_wrap_lowering)\nbatching.primitive_batchers[random_wrap_p] = random_wrap_batch_rule\n\n\ndef random_unwrap(keys):\n  if not dtypes.issubdtype(keys.dtype, dtypes.prng_key):\n    raise TypeError(f'random_unwrap takes key array operand, got {keys.dtype=}')\n  return random_unwrap_p.bind(keys)\n\nrandom_unwrap_p = core.Primitive('random_unwrap')\nad.defjvp_zero(random_unwrap_p)\nbatching.defvectorized(random_unwrap_p)\n\n@random_unwrap_p.def_abstract_eval\ndef random_unwrap_abstract_eval(keys_aval):\n  return core.physical_aval(keys_aval)\n\n@random_unwrap_p.def_impl\ndef random_unwrap_impl(keys):\n  return keys._base_array\n\ndef random_unwrap_lowering(ctx, keys):\n  return [keys]\n\nmlir.register_lowering(random_unwrap_p, random_unwrap_lowering)","sourceCodeStart":759,"sourceCodeEnd":795,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/prng.py#L759-L795","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Only unwrap typed keys: guard with jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key) before calling","If you already have raw data, skip the unwrap","Use the public jax.random.key_data instead of internal random_unwrap"],"exampleFix":"// before\nraw = jax.random.key_data(already_raw_uint32_array)\n\n// after\nimport jax, jax.numpy as jnp\nraw = (jax.random.key_data(x)\n       if jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key)\n       else jnp.asarray(x, dtype=jnp.uint32))","handlingStrategy":"type-guard","validationCode":"import jax\nif jax.dtypes.issubdtype(keys.dtype, jax.dtypes.prng_key):\n    raw = jax.random.key_data(keys)\nelse:\n    raw = keys  # already unwrapped","typeGuard":"import jax\ndef is_typed_key(x) -> bool:\n    return jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key)","tryCatchPattern":null,"preventionTips":["Guard unwrap calls with a prng_key dtype check","Avoid double-unwrap: track whether data is wrapped or raw"],"tags":["jax","prng","unwrap","dtype"],"backgroundTag":"jax-prng-key-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}