{"record":{"id":"05d88b15677129b1","repo":"jax-ml/jax","slug":"prng-key-seed-must-be-a-scalar-got-seed-r-05d88b","errorCode":null,"errorMessage":"PRNG key seed must be a scalar; got {seed!r}.","messagePattern":"PRNG key seed must be a scalar; got (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/random/threefry2x32.py","lineNumber":64,"sourceCode":"\ndef threefry_seed(seed: typing.Array) -> typing.Array:\n  \"\"\"Create a single raw threefry PRNG key from an integer seed.\n\n  Args:\n    seed: a 64- or 32-bit integer used as the value of the key.\n\n  Returns:\n    The PRNG key contents, modeled as an array of shape (2,) and dtype\n    uint32. The key is constructed from a 64-bit seed by effectively\n    bit-casting to a pair of uint32 values (or from a 32-bit seed by\n    first padding out with zeros).\n  \"\"\"\n  return _threefry_seed(seed)\n\n@api.jit(inline=True)\ndef _threefry_seed(seed: typing.Array) -> typing.Array:\n  if seed.shape:\n    raise TypeError(f\"PRNG key seed must be a scalar; got {seed!r}.\")\n  if not np.issubdtype(seed.dtype, np.integer):\n    raise TypeError(f\"PRNG key seed must be an integer; got {seed!r}\")\n  convert = lambda k: lax.expand_dims(lax.convert_element_type(k, np.uint32), [0])\n  k1 = convert(\n      lax.shift_right_logical(seed, lax._const(seed, 32)))\n  with config.numpy_dtype_promotion('standard'):\n    # TODO(jakevdp): in X64 mode, this can generate 64-bit computations for 32-bit\n    # inputs. We should avoid this.\n    k2 = convert(jnp.bitwise_and(seed, np.uint32(0xFFFFFFFF)))\n  return lax.concatenate([k1, k2], 0)\n\n\ndef _make_rotate_left(dtype):\n  if not dtypes.issubdtype(dtype, np.integer):\n    raise TypeError(\"_rotate_left only accepts integer dtypes.\")\n  nbits = np.array(dtypes.iinfo(dtype).bits, dtype)\n\n  def _rotate_left(x, d):","sourceCodeStart":46,"sourceCodeEnd":82,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/threefry2x32.py#L46-L82","documentation":"threefry2x32 seeds are converted to key material by splitting a scalar integer into two uint32 words. _threefry_seed (behind jax.random.PRNGKey and jax.random.key) raises this TypeError when the seed has a non-empty shape — a batch of integers cannot be a single seed, even if it has length 1.","triggerScenarios":"jax.random.PRNGKey(jnp.array([0])) or PRNGKey(np.array([1,2])); passing shape-(1,) arrays from config parsing; feeding the output of jnp.arange into PRNGKey expecting vectorized key creation.","commonSituations":"Config systems that wrap scalars in arrays; users expecting vectorized seeding (use jax.random.vmap or jax.vmap(jax.random.PRNGKey)(seeds) instead); test helpers parameterized over seeds-as-arrays.","solutions":["Pass a true scalar: jax.random.key(int(seed_array)) or .item() first","For many seeds, use jax.vmap(jax.random.PRNGKey)(jnp.array([1,2,3]))","Validate seed.ndim == 0 before calling in generic code"],"exampleFix":"// before\nseeds = jnp.array([1, 2, 3])\nkeys = jax.random.PRNGKey(seeds)  # TypeError\n\n// after\nkeys = jax.vmap(jax.random.PRNGKey)(jnp.array([1, 2, 3]))","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nseed = jnp.asarray(seed)\nassert seed.ndim == 0, f'seed must be scalar, got shape {seed.shape}'\nkey = jax.random.key(int(seed))","typeGuard":"import jax.numpy as jnp\ndef is_scalar_seed(s) -> bool:\n    return not jnp.asarray(s).shape","tryCatchPattern":null,"preventionTips":["Pass Python ints as seeds; call .item() on array seeds","Use jax.vmap(jax.random.PRNGKey)(seeds) for vectorized key creation"],"tags":["jax","prng","seed","scalar","shape"],"backgroundTag":"jax-invalid-seed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}