{"record":{"id":"af433862da6c6960","repo":"jax-ml/jax","slug":"prng-keys-must-be-loaded-from-smem-did-you-set-th","errorCode":null,"errorMessage":"PRNG keys must be loaded from SMEM. Did you set the memory space to MemorySpace.SMEM in the BlockSpec for the PRNG key input?","messagePattern":"PRNG keys must be loaded from SMEM\\. Did you set the memory space to MemorySpace\\.SMEM in the BlockSpec for the PRNG key input\\?","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2302,"sourceCode":"  if isinstance(ref_aval.memory_space, tpu_core.AccMemorySpace):\n    raise ValueError(\n        \"Loading from an accumulator is not supported. Use `matmul_pop` \"\n        \"instead, which will additionally zero out the accumulator.\"\n    )\n\n  ref_block_shape, *_ = ctx.block_shapes\n  ref, ref_block_shape = _transform_ref(\n      ref, ref_aval, ref_block_shape, prev_transforms\n  )\n  ref_type = ir.MemRefType(ref.type)\n  is_smem_load = str(ref_type.memory_space) == \"#tpu.memory_space<smem>\"\n  (aval_out,) = ctx.avals_out\n  if isinstance(aval_out.dtype, prng.KeyTy) and pl_random.is_pallas_impl(\n      aval_out.dtype._impl\n  ):\n    # TODO(justinfu): Merge this with standard extended dtype handling.\n    if not is_smem_load:\n      raise ValueError(\"PRNG keys must be loaded from SMEM. Did you set \"\n                       \"the memory space to MemorySpace.SMEM in the \"\n                       \"BlockSpec for the PRNG key input?\")\n    return _prng_key_load_lowering_rule(ctx, *args_flat, args_tree=args_tree)\n  if should_physicalize_dtype(aval_out.dtype):\n    # pyrefly: ignore[bad-argument-type]\n    physical_element_aval = jax_core.physical_element_aval(aval_out.dtype)\n    idx = cast(NDIndexer, idx)\n    if idx.int_indexer_shape:\n      raise NotImplementedError()\n    elt_slices = [\n        indexing.Slice(0, size) for size in physical_element_aval.shape]\n    idx = NDIndexer(\n        indices=idx.indices + tuple(elt_slices),\n        shape=idx.shape + physical_element_aval.shape,\n        int_indexer_shape=(),\n    )\n    physical_out_dtype = physical_element_aval.dtype\n    physical_out_shape = jax_core.physical_shape(","sourceCodeStart":2284,"sourceCodeEnd":2320,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2284-L2320","documentation":"Raised when loading PRNG keys in a TPU Pallas kernel where the key's BlockSpec memory space is not SMEM. The Pallas random implementation requires keys to reside in scalar memory (SMEM) because they are loaded one scalar at a time via memref loads. Loading keys from VMEM or ANY space is unsupported.","triggerScenarios":"Calling random_draw / pl.random with a key input whose BlockSpec (or pltpu.BlockSpec) does not set memory_space=tpu.MemorySpace.SMEM.","commonSituations":"Forgetting the memory_space argument in the BlockSpec for the keys argument when using pallas random RNG on TPU; defaults put blocks in VMEM.","solutions":["Set memory_space=pltpu.MemorySpace.SMEM on the BlockSpec for the PRNG key input","Verify you are using the pallas-specific random implementation (jax.pallas random key type) and its required layout"],"exampleFix":"# before\nkeys_spec = pltpu.BlockSpec((1, 2))\n# after\nkeys_spec = pltpu.BlockSpec((1, 2), memory_space=pltpu.MemorySpace.SMEM)","handlingStrategy":"validation","validationCode":"# Before launching, assert the keys BlockSpec uses SMEM\nassert keys_spec.memory_space is not None and 'SMEM' in str(keys_spec.memory_space), 'PRNG keys need memory_space=SMEM'","typeGuard":"def key_spec_is_smem(spec) -> bool:\n    ms = getattr(spec, 'memory_space', None)\n    return ms is not None and str(ms).endswith('SMEM')","tryCatchPattern":null,"preventionTips":["Make a helper make_key_spec() that always sets SMEM for RNG keys","Add a unit test that checks every kernel's BlockSpec memory spaces before compile"],"tags":["jax","pallas","tpu","prng","random","smem"],"backgroundTag":"pallas-prng-smem-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}