jax-ml/jax · error · ValueError

The Pallas TPU lowering currently supports in memory space…

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.

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

• Expected an accumulator ref, got
• Gather only supports loading from VMEM, got
• Get does not support loading from
• Invalid memory space
• Loads and stores are only allowed on VMEM and SMEM…

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

Appendix: 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)