{"record":{"id":"9ac3b72ed7010ac5","repo":"jax-ml/jax","slug":"expected-counter-to-be-a-scalar-integer-ref-got","errorCode":null,"errorMessage":"Expected counter to be a scalar integer ref; got {self._counter}","messagePattern":"Expected counter to be a scalar integer ref; got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/stateful_rng.py","lineNumber":81,"sourceCode":"  >>> rng\n  StatefulPRNG(_base_key=Array((), dtype=key<fry>) overlaying:\n  [ 0 42], _counter=Ref(0, dtype=int32, weak_type=True))\n  \"\"\"\n  _base_key: Array\n  _counter: core.Ref\n\n  def __post_init__(self):\n    if self._base_key is api_util.SENTINEL:\n      return\n    if not (isinstance(self._base_key, Array)\n            and dtypes.issubdtype(self._base_key.dtype, dtypes.prng_key)):\n      raise ValueError(f\"Expected base_key to be a typed PRNG key; got {self._base_key}\")\n\n    # TODO(jakevdp): how to validate a traced mutable array?\n    if not (isinstance(self._counter, core.Ref) or\n            (isinstance(self._counter, core.Tracer)\n             and isinstance(self._counter.aval, state_types.AbstractRef))):\n      raise ValueError(f\"Expected counter to be a scalar integer ref; got {self._counter}\")\n\n  def key(self, shape: int | Sequence[int] = ()) -> Array:\n    \"\"\"Generate a new JAX PRNGKey, updating the internal state.\n\n    Args:\n      shape: an optional shape if returning multiple keys.\n\n    Returns:\n      A new, independent PRNG key with the same impl/dtype as\n      ``self._base_key``.\n\n    Examples:\n      >>> from jax.experimental import random\n      >>> rng = random.stateful_rng(0)\n      >>> rng.key()\n      Array((), dtype=key<fry>) overlaying:\n      [1797259609 2579123966]\n      >>> rng.key()","sourceCodeStart":63,"sourceCodeEnd":99,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/stateful_rng.py#L63-L99","documentation":"The stateful RNG in jax.experimental.random tracks draw counts in a counter that must be a mutable scalar integer ref (core.Ref, or a tracer with an AbstractRef aval), enabling safe state updates under jit. __post_init__ raises this ValueError when counter is a plain array, Python int, or anything else non-ref.","triggerScenarios":"Manually constructing the stateful-RNG state with counter=np.int64(0) or a jnp scalar; passing a regular JAX array as counter; refactoring internals where the counter object loses its Ref type.","commonSituations":"Users poking at the experimental stateful API internals instead of the public constructor; partial checkpointing/restoration of state; version changes in the experimental ref-state API surface.","solutions":["Do not construct the state object directly — call stateful_rng(seed=...) which builds a valid counter internally","If restoring state, re-derive the counter via the public API rather than serializing it","Check jax version compatibility for the experimental stateful module"],"exampleFix":"// before\nrng = stateful_rng(seed=0)\nrng._counter = jnp.int32(0)  # later use raises\n\n// after\nrng = stateful_rng(seed=0)  # keep internal counter untouched","handlingStrategy":"validation","validationCode":"# Users cannot easily construct a valid Ref counter; rely on the public API\nrng = stateful_rng(seed=0)  # constructs a well-formed counter internally","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never set or replace the internal _counter attribute","Treat the stateful RNG's state as opaque; rebuild via stateful_rng()"],"tags":["jax","prng","stateful-rng","ref","experimental"],"backgroundTag":"jax-stateful-rng-counter-invalid","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}