{"record":{"id":"f02ed03a370e2d7f","repo":"jax-ml/jax","slug":"can-only-load-a-single-key-per-load","errorCode":null,"errorMessage":"Can only load a single key per load.","messagePattern":"Can only load a single key per load\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2402,"sourceCode":"  )\n  prev_transforms, idx = _canonicalize_transforms_to_indexer(\n      ref_aval, transforms, transforms_avals\n  )\n  (aval_out,) = ctx.avals_out\n  assert isinstance(aval_out.dtype, prng.KeyTy)\n  key_shape = aval_out.dtype._impl.key_shape\n  ref_block_shape, *_ = ctx.block_shapes\n  idx = cast(NDIndexer, idx)\n  ref, ref_block_shape = _transform_ref(\n      ref, ref_aval, ref_block_shape, prev_transforms\n  )\n\n  if len(key_shape) != 2:\n    raise NotImplementedError(\"Seed key_data must be 1D.\")\n  if key_shape[0] != 1:\n    raise NotImplementedError(\"Leading dimension of seed key_data must be 1.\")\n  if not all(s == 1 for s in idx.shape):\n    raise NotImplementedError(\"Can only load a single key per load.\")\n  assert ref_block_shape[-2:] == key_shape, f\"{ref_block_shape=} {key_shape=}\"\n\n  load_ops = []\n  for i in range(key_shape[1]):\n    ref_shape = tuple(\n        dim for dim in ref_block_shape if dim is not pallas_core.squeezed\n    )\n    scalar_idx = NDIndexer(\n        indices=(*idx.indices, 0, i), shape=ref_shape, int_indexer_shape=()\n    )\n    starts, _, _, _, _ = _indexer_to_start_size_stride(\n        scalar_idx,\n        ref_block_shape,\n        cast_to_index=True,\n    )\n    load_ops.append(memref.load(ref, starts))\n  return KeyScalarBundle(scalars=load_ops, key_shape=tuple(key_shape))\n","sourceCodeStart":2384,"sourceCodeEnd":2420,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2384-L2420","documentation":"Raised when the index used to load PRNG keys selects more than one element, i.e. the indexer shape is not all 1s. The PRNG key load path is scalar-only: each load fetches exactly one key.","triggerScenarios":"pl.load on an SMEM PRNG key ref with a slice/index whose shape contains dims > 1, requesting multiple keys in one load.","commonSituations":"Slicing the key block with a range (e.g. keys[0:4]) hoping to get several keys at once; vectorizing RNG inside the kernel.","solutions":["Load one key at a time (scalar indexer) and loop or split keys inside the kernel","Use random_split/random_fold_in to derive multiple keys from the single loaded seed"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"# in-kernel: load keys[0:1, 0:1]-style scalar indexers only\nassert all(s == 1 for s in idx_shape), 'load one key at a time'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Load keys with scalar indexers, loop over draws","Use random_fold_in to vary keys rather than slicing key arrays"],"tags":["jax","pallas","tpu","prng","indexing"],"backgroundTag":"pallas-prng-key-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}