jax-ml/jax · error · TypeError

PRNG key seed must be an integer; got {seed!r}

Error message

PRNG key seed must be an integer; got {seed!r}

What it means

When constructing a Philox 2x32 PRNG key, the seed must have an integer dtype (np.issubdtype(seed.dtype, np.integer)). Float, bool, or complex seeds raise this TypeError because the seed is bit-split into uint32 halves.

Source

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

  """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. Cast to an integer before use: int(seed) or np.uint64(seed)
2. Use jax.random.PRNGKey with an int (the top-level API also accepts uint64 shapes, but keep it integral)
3. Store seeds as int in configs to avoid silent rounding when casting floats

Example fix

# before
key = jax.random.PRNGKey(42.0)
# after
key = jax.random.PRNGKey(int(42.0))

import numpy as np
seed = int(seed) if not np.issubdtype(np.asarray(seed).dtype, np.integer) else seed
key = jax.random.PRNGKey(seed)

def is_integer_seed(seed) -> bool:
    import numpy as np
    return np.issubdtype(np.asarray(seed).dtype, np.integer)

When it happens

Trigger: jax.random.PRNGKey(42.0) or PRNGKey(np.float32(7)) with the philox2x32 implementation; passing a Python bool.

Common situations: Seeds loaded from JSON/config as floats; time-based seeds like time.time() passed directly; bool seeds from flag variables.

Related errors

• 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 a scalar; got {seed!r}.
• PRNG key seed must be a scalar; 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/ccaab44e822c2523. Report an issue: GitHub.