{"record":{"id":"2fc18830f4c529c4","repo":"jax-ml/jax","slug":"leading-dimension-of-seed-key-data-must-be-1","errorCode":null,"errorMessage":"Leading dimension of seed key_data must be 1.","messagePattern":"Leading dimension of seed key_data must be 1\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2400,"sourceCode":"  ref_aval, transforms_avals, _, _ = args_tree.unflatten(\n      ctx.avals_in\n  )\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))","sourceCodeStart":2382,"sourceCodeEnd":2418,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2382-L2418","documentation":"Raised when loading PRNG seed key_data whose leading dimension is not 1. The TPU lowering expects exactly one key per load with layout (1, num_elems); a batched leading dimension is unimplemented.","triggerScenarios":"pl.load of a PRNG key ref where key_shape[0] != 1, e.g. loading a (k, 2) array of keys as the seed.","commonSituations":"Trying to load multiple keys at once for per-lane RNG instead of splitting keys inside the kernel with random_fold_in/split.","solutions":["Load a single key of shape (1, 2) and derive per-use keys in-kernel (split/fold_in)"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"# host side: ensure the key argument is a single key, not a batch\nassert key.ndim == 0 or key.shape == (1,), 'load one key of shape (1, 2), not a batch'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never pass arrays of keys to Pallas kernels","Derive subkeys with random_split inside the kernel"],"tags":["jax","pallas","tpu","prng","shape-validation"],"backgroundTag":"pallas-prng-key-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}