jax-ml/jax · error · ValueError

The Pallas TPU lowering currently requires that the last two

Error message

The Pallas TPU lowering currently requires that the last two dimensions of your block shape are divisible by 8 and 128 respectively, or be equal to the respective dimensions of the overall array. {extra_msg}{err_details}

What it means

On TPU, the vector/VLM units require the last two dimensions of each block shape to be divisible by 8 and 128 respectively, unless they equal the corresponding full-array dimensions. The Pallas TPU lowering enforces this hardware alignment constraint before emitting the pipelined Mosaic module, and raises this ValueError (with extra guidance for dynamic shape export) when violated.

Source

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

      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 = ""
        if pallas_core.dynamic_shapes_export_enabled():
          extra_msg = (
              " In dynamic shape export - your kernel symbolic args must be"
              " annotated with constraints where the computation *after*"
              " applying any grid mapping is divisible by 8 and 128"
              " respectively. Ex: (mod(floordiv(m_dim, grid_size), 8) == 0))"
          )
        raise ValueError(
            "The Pallas TPU lowering currently requires that the last two "
            "dimensions of your block shape are divisible by 8 and 128 "
            "respectively, or be equal to the respective dimensions of the "
            "overall array. "
            + extra_msg
            + err_details()
        )
    else:
      assert rank == 1
      if bm.array_aval.dtype == jnp.bool_:
        bitwidth = dtypes.itemsize_bits(BOOL_MEMREF_TYPE)
      else:
        bitwidth = dtypes.itemsize_bits(physical_dtype)
      packing = 32 // bitwidth
      sublane_count = tpu_info.get_tpu_info().num_sublanes
      lane_count = tpu_info.get_tpu_info().num_lanes
      min_tiling = lane_count * packing
      chunk_size = sublane_count * lane_count

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Change the last dimension of your block shape to a multiple of 128 (e.g. 128, 256) or to the full array dimension
Make the second-to-last block dimension a multiple of 8 or equal to the array's dimension
For dynamic shape export, add divisibility constraints to symbolic args, e.g. (mod(floordiv(m_dim, grid_size), 8) == 0) and ... 128 == 0
Pad the input arrays so the block constraints can be satisfied and slice the result afterwards

Example fix

# before
BlockSpec(block_shape=(1024, 100), index_map=...)

# after
BlockSpec(block_shape=(1024, 128), index_map=...)  # last dim multiple of 128

Defensive patterns

Strategy: validation

Validate before calling

def check_tpu_block_alignment(block_shape, array_shape):
    ok_last = block_shape[-1] % 128 == 0 or block_shape[-1] == array_shape[-1]
    ok_prev = block_shape[-2] % 8 == 0 or block_shape[-2] == array_shape[-2]
    if not (ok_last and ok_prev):
        raise ValueError('last two block dims must be div by 128/8 or equal array dims')

Type guard

def is_tpu_aligned(block, arr) -> bool:
    return (block[-1] % 128 == 0 or block[-1] == arr[-1]) and \
           (block[-2] % 8 == 0 or block[-2] == arr[-2])

Try / catch

try:
    pallas_call(kernel, out_specs, in_specs, grid=grid)
except ValueError as e:
    if 'divisible by 8 and 128' in str(e):
        block = next_power_of_2_or_128(block)  # adjust tile and retry

Prevention

Default to tile sizes whose last dim is a multiple of 128 and second-to-last a multiple of 8
For dynamic shape export, add divisibility constraints on symbolic dims
Pad inputs to friendly sizes and slice outputs

When it happens

Trigger: A pallas_call on TPU whose BlockSpec block_shape has a last dim not divisible by 128 (and not equal to the array's last dim), or second-to-last dim not divisible by 8; with dynamic shape export, symbolic dims lacking constraints like (mod(floordiv(m_dim, grid_size), 8) == 0).

Common situations: Porting GPU Pallas/triton kernels to TPU with arbitrary tile sizes (e.g. 100x100 blocks); using non-multiple-of-128 tile sizes for the fastest dimension; dynamic shape export where symbolic dimensions aren't annotated with divisibility constraints.

Related errors

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