{"record":{"id":"35a62a6a471d2441","repo":"jax-ml/jax","slug":"when-used-within-transformed-code-jax-experimenta","errorCode":null,"errorMessage":"When used within transformed code, jax.experimental.random.stateful_rng() requires an explicit seed to be set.","messagePattern":"When used within transformed code, jax\\.experimental\\.random\\.stateful_rng\\(\\) requires an explicit seed to be set\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/stateful_rng.py","lineNumber":292,"sourceCode":"    >>> import jax\n    >>> jit_uniform = jax.jit(rng.uniform)\n    >>> jit_uniform()\n    Array(0.6672406, dtype=float32)\n    >>> jit_uniform()\n    Array(0.3890121, dtype=float32)\n\n    Keys can be generated directly if desired:\n\n    >>> rng.key()\n    Array((), dtype=key<fry>) overlaying:\n    [2954079971 3276725750]\n    >>> rng.key()\n    Array((), dtype=key<fry>) overlaying:\n    [2765691542  824333390]\n  \"\"\"\n  if seed is None:\n    if not core.trace_ctx.is_top_level():\n      raise TypeError(\n        \"When used within transformed code, jax.experimental.random.stateful_rng()\"\n        \" requires an explicit seed to be set.\")\n    entropy = np.random.SeedSequence().entropy\n    assert isinstance(entropy, int)\n    seed = np.int64(entropy & np.iinfo(np.int64).max)\n  assert seed is not None\n  return StatefulPRNG(\n    _base_key=random.key(seed, impl=impl),\n    _counter=ref.new_ref(0)\n  )\n","sourceCodeStart":274,"sourceCodeEnd":303,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/stateful_rng.py#L274-L303","documentation":"stateful_rng() with no seed normally derives a random seed from OS entropy at top level. Inside a JAX transformation (jit, grad, vmap, pmap, scan, etc.) that is disallowed — trace-time state creation with implicit entropy would break tracing and reproducibility — so it raises this TypeError unless seed is passed explicitly.","triggerScenarios":"Calling stateful_rng() inside a @jax.jit function, in a grad/vmap/scan body, or during any non-top-level trace; constructing generators lazily inside model init functions that get jitted.","commonSituations":"Model constructors run under hk.transform/jit/Flax NNX where the code creates an RNG without threading a seed in; refactoring top-level generator creation into helper functions called from traced code.","solutions":["Create the generator outside the transformation and pass it in as an argument or closure","Pass an explicit seed: stateful_rng(seed=some_int) when creation must happen inside","Use spawn()/split() on an outer generator to derive child keys inside traced code"],"exampleFix":"// before\n@jax.jit\ndef f(x):\n    rng = stateful_rng()  # TypeError inside transformed code\n    return x + rng.random()\n\n// after\n@jax.jit\ndef f(x, rng):\n    return x + rng.random()\nrng = stateful_rng(seed=0)\ny = f(x, rng)","handlingStrategy":"validation","validationCode":"import jax\ndef make_rng(seed=None):\n    if seed is None and not jax.core.trace_ctx.is_top_level():\n        raise ValueError('pass an explicit seed inside transformed code')\n    return stateful_rng(seed)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Create generators at top level and thread them as arguments","Always pass seed= when construction happens inside jit/scan/vmap"],"tags":["jax","prng","stateful-rng","jit","tracing"],"backgroundTag":"jax-traced-side-effect-not-allowed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}