{"record":{"id":"6f3ae36be425dbd2","repo":"jax-ml/jax","slug":"the-pallas-tpu-lowering-currently-supports-in-memo","errorCode":null,"errorMessage":"The Pallas TPU lowering currently supports in memory space ANY only blocks having the same block shape as the array shape and a trivial index_map (returning all 0s).{err_details}","messagePattern":"The Pallas TPU lowering currently supports in memory space ANY only blocks having the same block shape as the array shape and a trivial index_map \\(returning all 0s\\)\\.(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":986,"sourceCode":"      continue\n\n    def err_details():\n      return (f\"Block spec for {bm.origin} in pallas_call {debug_info.func_src_info} \"\n              \"has block shape \"\n              f\"{physical_block_shape}, array shape {physical_array_shape}, \"\n              # TODO(necula): add index_map source location info\n              f\"and index_map {bm.index_map_jaxpr}, in \"\n              f\"memory space {bm.block_aval.memory_space!r}.\"\n              \"\\nSee details at https://docs.jax.dev/en/latest/pallas/grid_blockspec.html#pallas-blockspec\")\n    if rank < 1:\n      raise ValueError(\n          \"The Pallas TPU lowering currently supports only blocks of \"\n          \"rank >= 1. \" + err_details())\n\n    if (\n        memory_space is ANY or memory_space == tpu_core.MemorySpace.HBM\n    ) and not bm.has_trivial_window():\n      raise ValueError(\n          \"The Pallas TPU lowering currently supports in memory space ANY \"\n          \"only blocks having the same block shape as the array shape \"\n          \"and a trivial index_map (returning all 0s).\" + err_details())\n\n    unmapped_bs = pallas_core._get_block_shape(physical_block_shape)\n    bs0, as0 = unmapped_bs[-1], physical_array_shape[-1]\n    if rank >= 2:\n      bs1, as1 = unmapped_bs[-2], physical_array_shape[-2]\n    else:\n      bs1, as1 = 1, 1\n\n    if rank >= 2:\n      evenly_divisible = (\n          (bs0 == as0 or bs0 % 128 == 0) and\n          (bs1 == as1 or bs1 % 8 == 0)\n      )\n      if not evenly_divisible:\n        extra_msg = \"\"","sourceCodeStart":968,"sourceCodeEnd":1004,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L968-L1004","documentation":"This error is thrown by JAX's Pallas TPU lowering when a BlockMapping whose memory space is ANY (or HBM) has a non-trivial window. For memory space ANY the TPU lowering can only handle blocks that cover the whole array (block shape == array shape) with a trivial index_map returning all 0s, because the compiler cannot infer partial-window semantics for unspecified memory spaces.","triggerScenarios":"Calling pallas_call (or pallas.tpu_kernel) with a GridMapping where a BlockMapping has memory_space=pallas_core.MemorySpace.ANY (the default) and either a non-trivial window (start_index_map/block_shape differ from the array) or a non-trivial index map. Typically happens when the compiler autodiff/partitions a kernel into ANY-memory-space intermediate buffers with windowed access.","commonSituations":"Writing a Pallas kernel where an intermediate (compiler-introduced) buffer gets memory space ANY but is accessed with a strided/offset window; using BlockSpec with non-trivial block_shape on TPU without explicitly setting memory_space; kernels that worked on GPU (where ANY is fine) ported to TPU.","solutions":["Set an explicit memory_space (e.g. pallas.MemorySpace.TPU_HBM or DMEM) on the block mapping so the ANY path is not taken","Make the block shape equal to the full array shape and use a trivial (all-zeros) index_map for ANY-memory-space operands","Restructure the kernel so ANY-memory-space operands are only accessed with trivial windows","Check that autodiff/transpose of your kernel isn't introducing windowed accesses on ANY buffers; mark those buffers explicitly instead"],"exampleFix":"# before\nBlockSpec(index_map=lambda i: i, block_shape=block)  # memory_space defaults to ANY\n\n# after\nBlockSpec(index_map=lambda i: i, block_shape=block,\n          memory_space=pallas.MemorySpace.TPU_HBM)","handlingStrategy":"validation","validationCode":"from jax._src.pallas import pallas_core\nfrom jax._src.pallas.mosaic import tpu_core\nfor bm in grid_mapping.block_mappings:\n    ms = bm.memory_space\n    if (ms is pallas_core.MemorySpace.ANY or ms == tpu_core.MemorySpace.HBM) and not bm.has_trivial_window():\n        raise ValueError(f'ANY/HBM block {bm} needs a trivial window')","typeGuard":"def has_only_trivial_any_windows(gm) -> bool:\n    return all(bm.has_trivial_window() for bm in gm.block_mappings\n               if bm.memory_space is pallas_core.MemorySpace.ANY)","tryCatchPattern":"try:\n    compiled = pallas_call(...)\nexcept ValueError as e:\n    if 'trivial index_map' in str(e):\n        # set explicit memory_space on BlockSpecs and retry\n        ...","preventionTips":["Always set an explicit memory_space on BlockSpecs used on TPU","Keep ANY-memory-space operands full-array with all-zero index maps","Test kernels with a tiny grid in CI to catch mapping validation early"],"tags":["jax","pallas","tpu","block-mapping","memory-space"],"backgroundTag":"pallas-tpu-unsupported-block-mapping","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}