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

philox4x32_seed requires a scalar seed; if seed.shape is non-empty it raises this TypeError with the offending value. This is the key-construction path for the Philox 4x32 PRNG implementation.

Source

Thrown at jax/_src/random/philox4x32.py:152

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


def philox4x32_seed(seed: typing.Array) -> typing.Array:
  """Create a single Philox 4x32 PRNG key from an integer seed."""
  return _philox4x32_seed(seed)


@api.jit(inline=True)
def _philox4x32_seed(seed: typing.Array) -> typing.Array:
  """Internal implementation of philox4x32_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 philox4x32 to mix the seed bits into a 2-word key.
  # Use the seed halves as counter words so both influence the output.
  out = philox4x32_p.bind(
      np.uint32(0),
      np.uint32(0),
      k0,
      k1,
      np.uint32(0),
      np.uint32(0),
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a scalar: PRNG seed of Python int or 0-d array
2. Use jax.vmap(philox4x32_seed) or jax.random.split for many keys
3. Squeeze accidental extra dims: seed.squeeze() before use

Example fix

# before
seeds = np.array([1, 2, 3])
keys = philox4x32_seed(seeds)
# after
keys = jax.vmap(philox4x32_seed)(seeds)

import numpy as np
assert np.asarray(seed).ndim == 0, 'philox4x32 seed must be scalar'

def is_scalar_seed(seed) -> bool:
    import numpy as np
    return np.asarray(seed).ndim == 0

When it happens

Trigger: Passing an array seed (e.g. np.array([1,2,3])) or a shape-(1,) array when creating a philox4x32 key.

Common situations: Vectorized per-sample seeds passed in one call instead of vmap; seeds unpacked from data files as arrays.

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/34053e2931f06224. Report an issue: GitHub.