jax-ml/jax · error · ValueError

Kernel input {j} in HBM but does not have trivial BlockSpec.

Error message

Kernel input {j} in HBM but does not have trivial BlockSpec.

What it means

Inside the TPU interpret-mode grid loop, an input that lives in HBM must have a trivial BlockSpec: its block shape must equal the full input shape (the whole tensor is passed to the kernel each iteration). If the interpreter sees an HBM input whose kernel-visible block shape differs from the input shape, it raises this ValueError naming the input index.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:2368

          token = callback.io_callback(
              # TODO(jburnim): Pass source_info from the pallas_call, in case this
              # store is involved in a data race.
              store,
              TOKEN_SHAPE_DTYPE,
              token,
              device_id,
              core_index,
              TPU_MEMORY_SPACE_IDXS[memory_space],
              input_ids[index],
              (),
              sliced_val,
          )
          return token

        for j, var in enumerate(input_vars):
          if input_var_memory_spaces[j] is _HBM:
            if var.aval.shape != block_shapes[j]:
              raise ValueError(
                  f'Kernel input {j} in HBM but does not have trivial'
                  ' BlockSpec.'
              )
            continue
          assert len(cur_start_indices[j].shape) == 1
          assert len(prev_start_indices[j].shape) == 1
          token = jax.lax.cond(
              (iteration_idx == initial_iteration_idx)
              | jax.lax.reduce_or(
                  cur_start_indices[j] != prev_start_indices[j], axes=(0,)
              ),
              functools.partial(
                  _store_slice_to_kernel_input,
                  j,
                  var,
                  input_var_memory_spaces[j],
              ),
              lambda t: t,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Give HBM inputs a BlockSpec with block shape equal to the full tensor shape (identity mapping), or None for scalars
Keep HBM inputs as copy sources (async_copy) rather than block-addressed arguments
Verify each input's memory space annotation matches its BlockSpec granularity

Example fix

# before
hbm_spec = BlockSpec((BM, BN), lambda i, j: (i*BM, j*BN))  # blocked HBM input
# after
hbm_spec = None if scalar else BlockSpec((M, N), lambda i, j: (0, 0))  # whole tensor, trivial

Defensive patterns

Strategy: validation

Validate before calling

for j, (x, spec) in enumerate(zip(inputs, in_specs)):
    if memory_spaces[j] == 'hbm' and spec is not None:
        assert tuple(spec.block_shape) == tuple(x.shape), f'input {j} in HBM needs trivial BlockSpec'

Type guard

def has_trivial_blockspec(spec, shape) -> bool:
    return spec is None or tuple(spec.block_shape) == tuple(shape)

Try / catch

try:
    interpret_run(kernel)
except ValueError as e:
    if 'does not have trivial BlockSpec' in str(e) and 'input' in str(e):
        # replace input spec with whole-tensor trivial spec and retry
        raise

Prevention

Give HBM args whole-tensor BlockSpecs or None
Stage HBM blocks via async_copy instead of blocked args
Review memory-space vs BlockSpec granularity per input

When it happens

Trigger: Passing an HBM input (e.g., via memory space annotations or mosaic params) with a BlockSpec that slices it into blocks, e.g. BlockSpec((BM,), ...) for an input of shape (N, BM).

Common situations: Marking an input as HBM while reusing the same blocked BlockSpec as VMEM inputs; kernels that stream blocks expecting HBM refs with per-iteration block views, which TPU interpret mode does not support.

Related errors

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