jax-ml/jax · error · ValueError

PRNG with name {impl.name} already registered: {impl}

Error message

PRNG with name {impl.name} already registered: {impl}

What it means

JAX keeps a global registry prngs mapping implementation names to PRNGImpl objects. register_prng raises ValueError if an implementation with the same name is already registered, preventing silent replacement of built-in impls like 'threefry2x32' or 'philox4x32'.

Source

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

  def __hash__(self) -> int:
    return hash(self.tag)

  def __str__(self) -> str:
    return self.tag

  def pprint(self):
    ty = self.__class__.__name__
    return (pp.text(f"{ty} [{self.tag}] {{{self.name}}}:") +
            pp.nest(2, pp.group(pp.brk() + pp.join(pp.brk(), [
              pp.text(f"{k} = {v}") for k, v in self._asdict().items()
            ]))))


prngs: dict[str, PRNGImpl] = {}

def register_prng(impl: PRNGImpl):
  if impl.name in prngs:
    raise ValueError(f'PRNG with name {impl.name} already registered: {impl}')
  prngs[impl.name] = impl


# -- PRNG key arrays

def _check_prng_key_data(impl, key_data: typing.Array):
  ndim = len(impl.key_shape)
  if not all(hasattr(key_data, attr) for attr in ['ndim', 'shape', 'dtype']):
    raise TypeError("JAX encountered invalid PRNG key data: expected key_data "
                    f"to have ndim, shape, and dtype attributes. Got {key_data}")
  if key_data.ndim < 1:
    raise TypeError("JAX encountered invalid PRNG key data: expected "
                    f"key_data.ndim >= 1; got ndim={key_data.ndim}")
  if key_data.shape[-ndim:] != impl.key_shape:
    raise TypeError("JAX encountered invalid PRNG key data: expected key_data.shape to "
                    f"end with {impl.key_shape}; got shape={key_data.shape} for {impl=}")
  if key_data.dtype not in [np.uint32, float0]:
    raise TypeError("JAX encountered invalid PRNG key data: expected key_data.dtype = uint32; "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard registration: if impl.name not in prngs: register_prng(impl)
2. Give the custom implementation a unique name
3. Reuse the already-registered impl instead of re-registering (fetch prngs[impl.name] and compare)

Example fix

# before
register_prng(my_impl)  # raises if re-imported
# after
from jax._src.random import prng
if my_impl.name not in prng.prngs:
    prng.register_prng(my_impl)

from jax._src.random import prng
if my_impl.name not in prng.prngs:
    prng.register_prng(my_impl)

def is_new_prng_name(name) -> bool:
    from jax._src.random import prng
    return name not in prng.prngs

try:
    register_prng(my_impl)
except ValueError:
    pass  # already registered; reuse existing

When it happens

Trigger: Calling jax._src.random.prng.register_prng with a custom PRNGImpl whose name collides with an existing one, or registering the same custom impl twice (e.g. module re-import or repeated notebook cell execution).

Common situations: Custom PRNG plugins in libraries that register on import; notebooks that re-run a registration cell; library version upgrades adding a name your custom impl also uses.

Related errors

• Duplicate serialization registration for type `{nodetype}`.
• the rematted computation's closure contains a mutable array
• Effects not supported in partial-eval of `checkpoint`/`remat
• Config option {name} already defined
• Duplicate serialization registration for serialized_name `{s

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