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

philox4x32_seed requires the seed to have an integer dtype; float, complex, or bool seeds raise this TypeError because the seed is decomposed into 32-bit integer words.

Source

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

  """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),
  )
  return jnp.array([out[0], out[1]], dtype=np.uint32)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to int first: int(seed), np.int64(seed), or jnp.asarray(seed, jnp.uint32)
2. Canonicalize seed handling in one helper that coerces to int

Example fix

# before
key = philox4x32_seed(jnp.array(1.5))
# after
key = philox4x32_seed(jnp.array(1, dtype=np.uint32))

import numpy as np
if not np.issubdtype(np.asarray(seed).dtype, np.integer):
    seed = int(seed)

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

When it happens

Trigger: Calling philox4x32_seed(jnp.float32(0.0)) or PRNGKey with a float seed under the philox4x32 implementation.

Common situations: Float seeds from configuration files or time.time(); np.bool_ seeds from flags.

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