{"record":{"id":"89813669a5d37532","repo":"jax-ml/jax","slug":"prng-key-seed-must-be-a-scalar-got-seed-r","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/philox2x32.py","lineNumber":141,"sourceCode":"\ndef _is_philox2x32_key(key: typing.Array) -> bool:\n  \"\"\"Return True if the input is a valid Philox 2x32 PRNG key.\"\"\"\n  try:\n    return key.shape == (1,) and key.dtype == np.uint32\n  except AttributeError:\n    return False\n\n\ndef philox2x32_seed(seed: typing.Array) -> typing.Array:\n  \"\"\"Create a single Philox 2x32 PRNG key from an integer seed.\"\"\"\n  return _philox2x32_seed(seed)\n\n\n@api.jit(inline=True)\ndef _philox2x32_seed(seed: typing.Array) -> typing.Array:\n  \"\"\"Internal implementation of philox2x32_seed.\"\"\"\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.convert_element_type(k, np.uint32)\n  k0 = convert(\n      lax.shift_right_logical(seed, lax.convert_element_type(32, seed.dtype))\n  )\n  with config.numpy_dtype_promotion(\"standard\"):\n    k1 = convert(jnp.bitwise_and(seed, np.uint32(0xFFFFFFFF)))\n  # Hash through philox2x32 to mix the seed bits into a 1-word key.\n  # Use both seed halves as counter words so they both influence the output.\n  out0, _ = philox2x32_p.bind(np.uint32(0), k0, k1)\n  return jnp.array([out0], dtype=np.uint32)\n\n\ndef philox2x32_split(key: typing.Array, shape: prng.Shape) -> typing.Array:\n  \"\"\"Split a Philox 2x32 PRNG key into multiple sub-keys.\"\"\"\n  shape = tuple(map(core.concrete_dim_or_error, shape))\n  return _philox2x32_split(key, shape)","sourceCodeStart":123,"sourceCodeEnd":159,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/random/philox2x32.py#L123-L159","documentation":"philox2x32_seed (used to build a PRNG key from a raw seed) requires the seed to be a scalar array. If seed.shape is non-empty, a TypeError is raised with the offending value. This guards the key-construction path jax.random.PRNGKey(..., impl='philox2x32').","triggerScenarios":"Passing an array seed (e.g. jax.random.PRNGKey(np.array([1,2]))) or an implied key with the philox2x32 implementation; passing a shape-(1,) array like np.array([42]).","commonSituations":"Feeding a per-example seed vector for reproducible per-sample noise; converting old seeds stored as arrays; using the threefry2x32-style API where a length-2 seed was tolerated.","solutions":["Pass a scalar seed: jax.random.PRNGKey(42) or jnp.asarray(42, dtype=np.uint32)","For per-element seeds, use jax.vmap over scalar seeds or jax.random.split of one key","If a shape-(1,) array sneaks in, index it: seed[0]"],"exampleFix":"# before\nkey = jax.random.PRNGKey(np.array([42]), impl='philox2x32')\n# after\nkey = jax.random.PRNGKey(42, impl='philox2x32')","handlingStrategy":"validation","validationCode":"import numpy as np\nseed = np.asarray(seed)\nassert seed.ndim == 0, f'seed must be scalar, got shape {seed.shape}'\nkey = jax.random.PRNGKey(int(seed), impl='philox2x32')","typeGuard":"def is_scalar_int(seed) -> bool:\n    import numpy as np\n    a = np.asarray(seed)\n    return a.ndim == 0 and np.issubdtype(a.dtype, np.integer)","tryCatchPattern":null,"preventionTips":["Store seeds as Python ints","Use vmap over seed arrays for per-sample keys"],"tags":["jax","prng","seed","philox"],"backgroundTag":"prng-seed-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}