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

PRNG seeds must be integers because threefry hashes the integer bits into two uint32 key words. _threefry_seed raises this TypeError when the seed is a float (e.g. 42.0 or np.float64), bool-kind arrays are the only tolerated exception downstream in some paths, and non-integer inputs cannot be hashed deterministically across platforms.

Source

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

  """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):
    if lax.dtype(d) != dtype:
      d = lax.convert_element_type(d, dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Coerce to int: jax.random.key(int(seed))
2. Fix config parsing to yield ints (e.g. schema validation)
3. When computing seeds arithmetically, wrap the result in int()

Example fix

// before
key = jax.random.PRNGKey(42.0)

// after
key = jax.random.PRNGKey(42)

seed = int(seed) if not hasattr(seed, 'dtype') else seed
# or generically:
import numpy as np
if not np.issubdtype(np.asarray(seed).dtype, np.integer):
    seed = int(seed)

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

When it happens

Trigger: jax.random.PRNGKey(42.0); PRNGKey(np.float32(0)); seeds read from JSON/config as floats; dividing a seed expression ('seed/n_devices') producing floats.

Common situations: JSON/YAML configs where 0 parses as int but 0.0 or computed values parse as float; hyperparameter sweeps computing seeds arithmetically; np.random seeds passed through float dtypes.

Related errors

• mulhi requires integer inputs, got {dtype}
• 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}

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