{"record":{"id":"6872498fbb360190","repo":"jax-ml/jax","slug":"seed-key-data-must-be-1d","errorCode":null,"errorMessage":"Seed key_data must be 1D.","messagePattern":"Seed key_data must be 1D\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2398,"sourceCode":"  \"\"\"\n  ref, transforms, _, _ = args_tree.unflatten(args_flat)\n  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,","sourceCodeStart":2380,"sourceCodeEnd":2416,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2380-L2416","documentation":"Raised when the key_data of a PRNG seed loaded in a TPU Pallas kernel is not effectively 1D. The lowering requires key_data shape (1, N) after accounting for the leading impl dimension; anything else (e.g. a 2D key array) is unimplemented.","triggerScenarios":"pl.load on an SMEM PRNG key ref whose key_shape (from aval_out.shape plus impl layout) has len != 2, i.e. key data that isn't a flat 1D array of shape (1, n).","commonSituations":"Passing a shaped/random key array (e.g. shape (4, 2)) as the seed input rather than a single unraveled key; using a non-pallas RNG impl shape.","solutions":["Pass a single seed key with 1D key_data: use jax.random.key / pallas random seeding so key_data is (1, 2) or (1, n)","Generate multiple keys inside the kernel via random_draw from one seed instead of loading a key array"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"import jax\nkey = jax.random.key(seed)\nassert key.shape == ()  # single key; key_data stays 1D (1, n)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pass exactly one PRNG seed key per kernel input","Split/fold keys inside the kernel, not via the input's shape"],"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"}