jax-ml/jax · error · TypeError

threefry_2x32 requires uint32 arguments, got {}

Error message

threefry_2x32 requires uint32 arguments, got {}

What it means

threefry_2x32 is the low-level 2x32 block cipher at the heart of JAX's default PRNG; its kernel only accepts uint32 operands for both key words and the count/data array. This TypeError fires when any of key1, key2, or count has a different dtype, typically int32/int64 after implicit conversions.

Source

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

    platform='oneapi',
    inline=False)


@api.jit(inline=True)
def threefry_2x32(keypair, count):
  """Apply the Threefry 2x32 hash.

  Args:
    keypair: a pair of 32bit unsigned integers used for the key.
    count: an array of dtype uint32 used for the counts.

  Returns:
    An array of dtype uint32 with the same shape as `count`.
  """
  key1, key2 = keypair
  if not lax.dtype(key1) == lax.dtype(key2) == lax.dtype(count) == np.uint32:
    msg = "threefry_2x32 requires uint32 arguments, got {}"
    raise TypeError(msg.format([lax.dtype(x) for x in [key1, key2, count]]))

  flat_count = count.ravel()
  odd_size = flat_count.shape[0] % 2
  if core.is_constant_dim(odd_size):
    if odd_size:
      x = list(jnp.split(jnp.concatenate([flat_count, jnp.uint32([0])]), 2))
    else:
      x = list(jnp.split(flat_count, 2))
  else:
    # With symbolic shapes we cannot always tell statically if odd_size is true
    # or false, so we rewrite this without a conditional.
    flat_count_padded = jnp.concatenate([flat_count, jnp.uint32([0])])
    flat_count_padded_half_size = flat_count_padded.shape[0] // 2
    x = [
      lax_slicing.dynamic_slice(flat_count_padded, (0,),
                                (flat_count_padded_half_size,)),
      lax_slicing.dynamic_slice(flat_count_padded,
                                (flat_count_padded_half_size,),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast all three operands: key1.astype(jnp.uint32), key2.astype(jnp.uint32), count.astype(jnp.uint32)
2. Prefer the public APIs (jax.random.fold_in, key_data round-trips) which handle dtype canonicalization
3. Be explicit with dtype in literals: jnp.uint32([...]) instead of jnp.array([...])

Example fix

// before
out = threefry_2x32((k1, k2), jnp.arange(8))

// after
out = threefry_2x32((k1, k2), jnp.arange(8, dtype=jnp.uint32))

import jax.numpy as jnp
key1, key2, count = (jnp.asarray(x, dtype=jnp.uint32) for x in (key1, key2, count))

import jax.numpy as jnp
def all_uint32(*xs) -> bool:
    return all(jnp.asarray(x).dtype == jnp.uint32 for x in xs)

When it happens

Trigger: Calling jax._src.random.threefry2x32.threefry_2x32 directly with int64 arrays (x64 mode); passing counts created by jnp.arange (int32 default) without casting; mixing numpy int arrays with uint32 keys.

Common situations: X64-enabled environments where literals default to int64; researchers using the low-level threefry API directly for custom hash chains; data loaded from int-typed sources fed as 'count'.

Related errors

• JAX encountered invalid PRNG key data: expected key_data.dty
• requires 8-, 16-, 32- or 64-bit field width.
• random_unwrap takes key array operand, got {keys.dtype=}
• requires 8-, 16-, 32- or 64-bit field width.
• 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/74b2c6aaf8243af1. Report an issue: GitHub.