jax-ml/jax · error · TypeError

PRNG key seed must be a scalar; got {seed!r}.

Error message

PRNG key seed must be a scalar; got {seed!r}.

What it means

philox2x32_seed (used to build a PRNG key from a raw seed) requires the seed to be a scalar array. If seed.shape is non-empty, a TypeError is raised with the offending value. This guards the key-construction path jax.random.PRNGKey(..., impl='philox2x32').

Source

Thrown at jax/_src/random/philox2x32.py:141

def _is_philox2x32_key(key: typing.Array) -> bool:
  """Return True if the input is a valid Philox 2x32 PRNG key."""
  try:
    return key.shape == (1,) and key.dtype == np.uint32
  except AttributeError:
    return False


def philox2x32_seed(seed: typing.Array) -> typing.Array:
  """Create a single Philox 2x32 PRNG key from an integer seed."""
  return _philox2x32_seed(seed)


@api.jit(inline=True)
def _philox2x32_seed(seed: typing.Array) -> typing.Array:
  """Internal implementation of philox2x32_seed."""
  if seed.shape:
    raise TypeError(f"PRNG key seed must be a scalar; got {seed!r}.")
  if not np.issubdtype(seed.dtype, np.integer):
    raise TypeError(f"PRNG key seed must be an integer; got {seed!r}")
  convert = lambda k: lax.convert_element_type(k, np.uint32)
  k0 = convert(
      lax.shift_right_logical(seed, lax.convert_element_type(32, seed.dtype))
  )
  with config.numpy_dtype_promotion("standard"):
    k1 = convert(jnp.bitwise_and(seed, np.uint32(0xFFFFFFFF)))
  # Hash through philox2x32 to mix the seed bits into a 1-word key.
  # Use both seed halves as counter words so they both influence the output.
  out0, _ = philox2x32_p.bind(np.uint32(0), k0, k1)
  return jnp.array([out0], dtype=np.uint32)


def philox2x32_split(key: typing.Array, shape: prng.Shape) -> typing.Array:
  """Split a Philox 2x32 PRNG key into multiple sub-keys."""
  shape = tuple(map(core.concrete_dim_or_error, shape))
  return _philox2x32_split(key, shape)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a scalar seed: jax.random.PRNGKey(42) or jnp.asarray(42, dtype=np.uint32)
2. For per-element seeds, use jax.vmap over scalar seeds or jax.random.split of one key
3. If a shape-(1,) array sneaks in, index it: seed[0]

Example fix

# before
key = jax.random.PRNGKey(np.array([42]), impl='philox2x32')
# after
key = jax.random.PRNGKey(42, impl='philox2x32')

import numpy as np
seed = np.asarray(seed)
assert seed.ndim == 0, f'seed must be scalar, got shape {seed.shape}'
key = jax.random.PRNGKey(int(seed), impl='philox2x32')

def is_scalar_int(seed) -> bool:
    import numpy as np
    a = np.asarray(seed)
    return a.ndim == 0 and np.issubdtype(a.dtype, np.integer)

When it happens

Trigger: Passing an array seed (e.g. jax.random.PRNGKey(np.array([1,2]))) or an implied key with the philox2x32 implementation; passing a shape-(1,) array like np.array([42]).

Common situations: Feeding a per-example seed vector for reproducible per-sample noise; converting old seeds stored as arrays; using the threefry2x32-style API where a length-2 seed was tolerated.

Related errors

• PRNG key seed must be a scalar; got {seed!r}.
• PRNG key seed must be an integer; got {seed!r}
• philox2x32_random_bits got invalid prng key.
• requires 8-, 16-, 32- or 64-bit field width.
• random bits array of size exceeding 2 ** 64

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