jax-ml/jax · error · ValueError

The Pallas TPU lowering currently supports in memory space A

Error message

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}

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:986

      continue

    def err_details():
      return (f"Block spec for {bm.origin} in pallas_call {debug_info.func_src_info} "
              "has block shape "
              f"{physical_block_shape}, array shape {physical_array_shape}, "
              # TODO(necula): add index_map source location info
              f"and index_map {bm.index_map_jaxpr}, in "
              f"memory space {bm.block_aval.memory_space!r}."
              "\nSee details at https://docs.jax.dev/en/latest/pallas/grid_blockspec.html#pallas-blockspec")
    if rank < 1:
      raise ValueError(
          "The Pallas TPU lowering currently supports only blocks of "
          "rank >= 1. " + err_details())

    if (
        memory_space is ANY or memory_space == tpu_core.MemorySpace.HBM
    ) and not bm.has_trivial_window():
      raise ValueError(
          "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())

    unmapped_bs = pallas_core._get_block_shape(physical_block_shape)
    bs0, as0 = unmapped_bs[-1], physical_array_shape[-1]
    if rank >= 2:
      bs1, as1 = unmapped_bs[-2], physical_array_shape[-2]
    else:
      bs1, as1 = 1, 1

    if rank >= 2:
      evenly_divisible = (
          (bs0 == as0 or bs0 % 128 == 0) and
          (bs1 == as1 or bs1 % 8 == 0)
      )
      if not evenly_divisible:
        extra_msg = ""

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set an explicit memory_space (e.g. pallas.MemorySpace.TPU_HBM or DMEM) on the block mapping so the ANY path is not taken
2. Make the block shape equal to the full array shape and use a trivial (all-zeros) index_map for ANY-memory-space operands
3. Restructure the kernel so ANY-memory-space operands are only accessed with trivial windows
4. Check that autodiff/transpose of your kernel isn't introducing windowed accesses on ANY buffers; mark those buffers explicitly instead

Example fix

# before
BlockSpec(index_map=lambda i: i, block_shape=block)  # memory_space defaults to ANY

# after
BlockSpec(index_map=lambda i: i, block_shape=block,
          memory_space=pallas.MemorySpace.TPU_HBM)

from jax._src.pallas import pallas_core
from jax._src.pallas.mosaic import tpu_core
for bm in grid_mapping.block_mappings:
    ms = bm.memory_space
    if (ms is pallas_core.MemorySpace.ANY or ms == tpu_core.MemorySpace.HBM) and not bm.has_trivial_window():
        raise ValueError(f'ANY/HBM block {bm} needs a trivial window')

def has_only_trivial_any_windows(gm) -> bool:
    return all(bm.has_trivial_window() for bm in gm.block_mappings
               if bm.memory_space is pallas_core.MemorySpace.ANY)

try:
    compiled = pallas_call(...)
except ValueError as e:
    if 'trivial index_map' in str(e):
        # set explicit memory_space on BlockSpecs and retry
        ...

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Memory space {self.memory_space} is not supported by mesh {s
• Unsupported core type: {core_type}
• Invalid memory space: {memory_space!r}
• Loads are only allowed on VMEM and SMEM references.
• Loads and stores are only allowed on VMEM and SMEM reference

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/6f3ae36be425dbd2. Report an issue: GitHub.