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

threefry2x32 seeds are converted to key material by splitting a scalar integer into two uint32 words. _threefry_seed (behind jax.random.PRNGKey and jax.random.key) raises this TypeError when the seed has a non-empty shape — a batch of integers cannot be a single seed, even if it has length 1.

Source

Thrown at jax/_src/random/threefry2x32.py:64

def threefry_seed(seed: typing.Array) -> typing.Array:
  """Create a single raw threefry PRNG key from an integer seed.

  Args:
    seed: a 64- or 32-bit integer used as the value of the key.

  Returns:
    The PRNG key contents, modeled as an array of shape (2,) and dtype
    uint32. The key is constructed from a 64-bit seed by effectively
    bit-casting to a pair of uint32 values (or from a 32-bit seed by
    first padding out with zeros).
  """
  return _threefry_seed(seed)

@api.jit(inline=True)
def _threefry_seed(seed: typing.Array) -> typing.Array:
  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.expand_dims(lax.convert_element_type(k, np.uint32), [0])
  k1 = convert(
      lax.shift_right_logical(seed, lax._const(seed, 32)))
  with config.numpy_dtype_promotion('standard'):
    # TODO(jakevdp): in X64 mode, this can generate 64-bit computations for 32-bit
    # inputs. We should avoid this.
    k2 = convert(jnp.bitwise_and(seed, np.uint32(0xFFFFFFFF)))
  return lax.concatenate([k1, k2], 0)


def _make_rotate_left(dtype):
  if not dtypes.issubdtype(dtype, np.integer):
    raise TypeError("_rotate_left only accepts integer dtypes.")
  nbits = np.array(dtypes.iinfo(dtype).bits, dtype)

  def _rotate_left(x, d):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a true scalar: jax.random.key(int(seed_array)) or .item() first
2. For many seeds, use jax.vmap(jax.random.PRNGKey)(jnp.array([1,2,3]))
3. Validate seed.ndim == 0 before calling in generic code

Example fix

// before
seeds = jnp.array([1, 2, 3])
keys = jax.random.PRNGKey(seeds)  # TypeError

// after
keys = jax.vmap(jax.random.PRNGKey)(jnp.array([1, 2, 3]))

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

import jax.numpy as jnp
def is_scalar_seed(s) -> bool:
    return not jnp.asarray(s).shape

When it happens

Trigger: jax.random.PRNGKey(jnp.array([0])) or PRNGKey(np.array([1,2])); passing shape-(1,) arrays from config parsing; feeding the output of jnp.arange into PRNGKey expecting vectorized key creation.

Common situations: Config systems that wrap scalars in arrays; users expecting vectorized seeding (use jax.random.vmap or jax.vmap(jax.random.PRNGKey)(seeds) instead); test helpers parameterized over seeds-as-arrays.

Related errors

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

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