jax-ml/jax · error · ValueError

The Pallas TPU lowering currently requires that rank 1 block

Error message

The Pallas TPU lowering currently requires that rank 1 block shapes, either 1) the first (and only) dimension of the block shape is equal to the first (and only) dimension of the array shape, or 2) the first (and only) dimension of the block shape is a multiple of {chunk_size}, or 3) the first (and only) dimension of the block shape is a power of 2 and at least the tiling size ({min_tiling} = 128 * (32 // {dtypes.itemsize_bits(physical_dtype)})) of the array shape. {err_details}

What it means

For rank-1 arrays on TPU, the Pallas lowering can only pipeline certain block sizes: the block size must equal the array length, be a multiple of chunk_size, or be a power of 2 that is at least min_tiling = 128 * (32 // itemsize_bits(dtype)). This is because 1-D access must map efficiently onto the TPU's rectangular VLM layouts.

Source

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

        )
    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
      feasible_block_size = (
          bs0 == as0
          or bs0 % chunk_size == 0
          or (bs0 >= min_tiling and (bs0 & (bs0 - 1)) == 0)  # power of 2
      )
      if not feasible_block_size:
        raise ValueError(
            "The Pallas TPU lowering currently requires that rank 1 block"
            " shapes, either 1) the first (and only) dimension of the block"
            " shape is equal to the first (and only) dimension of the array"
            " shape, or 2) the first (and only) dimension of the block shape"
            f" is a multiple of {chunk_size}, or 3) the first (and only)"
            " dimension of the block shape is a power of 2 and at least the"
            f" tiling size ({min_tiling} = 128 * (32 //"
            f" {dtypes.itemsize_bits(physical_dtype)})) of the array shape. "
            + err_details()
        )


def lower_jaxpr_to_pipelined_module(
    lowering_context: mlir.LoweringRuleContext,
    grid_mapping: pallas_core.GridMapping,
    jaxpr: jax_core.Jaxpr,
    *,
    dimension_semantics: Sequence[tpu_core.DimensionSemantics] | None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use a power-of-2 block size at least the min_tiling for your dtype (e.g. >=128 for 32-bit types, >=512 for 8-bit types)
Use a block size that is a multiple of chunk_size (chunk_size = 128 * 128 // itemsize_bits)
Set the block size equal to the full array length (with a grid of 1) if it fits in memory
Switch the operand to 2-D by adding a trailing dimension of 1 if the algorithm permits

Example fix

# before
BlockSpec(block_shape=(100,), index_map=lambda i: i)  # 100 not feasible

# after
BlockSpec(block_shape=(128,), index_map=lambda i: i)  # power of 2 >= 128

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
def feasible_rank1(bs0, as0, dtype):
    import jax.numpy as jnp
    itemsize_bits = jnp.dtype(dtype).itemsize * 8
    chunk = 128 * 128 // itemsize_bits
    min_tiling = 128 * (32 // itemsize_bits)
    return bs0 == as0 or bs0 % chunk == 0 or (bs0 >= min_tiling and (bs0 & (bs0-1)) == 0)

Type guard

def is_valid_1d_block(bs0, as0, chunk_size, min_tiling) -> bool:
    return bs0 == as0 or bs0 % chunk_size == 0 or \
           (bs0 >= min_tiling and bs0 & (bs0 - 1) == 0)

Try / catch

try:
    pallas_call(kernel, out, grid=grid)
except ValueError as e:
    if 'rank 1 block' in str(e):
        out = pallas_call(kernel, reshape_2d_spec_out, grid=grid)  # reshape to 2D and retry

Prevention

For rank-1 kernels use power-of-2 blocks >= 128 (>= 512 for 8-bit dtypes)
Remember min_tiling doubles for 16-bit and quadruples for 8-bit dtypes
Prefer 2-D shapes with a trailing dim of 1 for 1-D algorithms

When it happens

Trigger: A rank-1 BlockSpec on TPU with a block size that is not equal to the array size, not a multiple of chunk_size, and not a power-of-2 >= min_tiling (e.g. a 100-element block of a 1000-element float32 array).

Common situations: 1-D kernels (e.g. elementwise maps over vectors, sorts, scans) tiled with odd block sizes; low-precision dtypes (f8/bf16) raising min_tiling to 512; block sizes chosen to match a dataset size rather than hardware constraints.

Related errors

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