jax-ml/jax · error · ValueError

Unsupported block dim type: {type(b)}

Error message

Unsupported block dim type: {type(b)}

What it means

When the interpreter computes block start indices from a BlockSpec's block mapping, each dimension's block type must be one of pallas_core.Element, Blocked, or an int. Any other type in the BlockSpec's block_shape raises this ValueError naming the offending type.

Source

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

  jaxpr = block_mapping.index_map_jaxpr
  token, block_indices = _interpret_jaxpr(
      jaxpr,
      *jaxpr.consts,
      *loop_idx,
      *args,
      ctx=ctx,
      token=token,
  )
  def _get_start_index(i, b):
    match b:
      case pallas_core.Squeezed():
        return i
      case pallas_core.Element():
        return i
      case pallas_core.Blocked():
        return i * b.block_size
      case _:
        raise ValueError(f"Unsupported block dim type: {type(b)}")
  ret = jnp.array(
      tuple(
          _get_start_index(i, b)
          for i, b in zip(block_indices, block_mapping.block_shape)
      ),
      dtype=jnp.int32,
  )
  return token, block_indices, ret


def _get_parallel_dim_semantics(
    mosaic_params: mosaic_core.CompilerParams, num_dimensions_in_grid: int,
) -> tuple[bool, ...]:
  """Returns a tuple indicating which grid dimensions have parallel semantics.

  Args:
    mosaic_params: The compiler params for the Mosaic TPU backend.
    num_dimensions_in_grid: The number of dimensions in the grid.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure every block_shape entry is an int, pallas_core.Element(), or pallas_core.Blocked(...)
2. Use None (no BlockSpec) for scalar/whole-tensor arguments rather than odd block types
3. Align JAX/jaxlib versions so pallas_core types match the interpreter

Example fix

# before
spec = BlockSpec((None, 128), lambda i: (0, i*128))
# after
from jax._src.pallas import pallas_core
spec = BlockSpec((pallas_core.Element(), 128), lambda i: (0, i*128))

from jax._src.pallas import pallas_core
valid = (int, pallas_core.Element, pallas_core.Blocked)
assert all(isinstance(b, valid) for b in spec.block_shape), f'bad block types: {[type(b) for b in spec.block_shape]}'

def is_valid_block_shape(bs) -> bool:
    from jax._src.pallas import pallas_core
    return all(isinstance(b, (int, pallas_core.Element, pallas_core.Blocked)) for b in bs)

try:
    interpret_run(kernel)
except ValueError as e:
    if 'Unsupported block dim type' in str(e):
        # normalize block_shape entries to int/Element/Blocked and retry
        raise

When it happens

Trigger: Constructing a BlockSpec whose block_shape entries are unexpected types (e.g., None, a string, a custom class, or numpy scalars in older versions) so the match statement in _get_start_index falls through to the error case.

Common situations: Passing None for broadcast/whole-tensor dims instead of pallas_core.Element; version mismatches between jax and jaxlib/pallas where BlockSpec internals changed; dynamically built block specs with heterogeneous entries.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Out-of-bounds block index {block_indices} for input "{input_
• Out-of-bounds block index {block_indices} for output "{outpu
• Kernel input {j} in HBM but does not have trivial BlockSpec.
• Kernel output {j} in HBM but does not have trivial BlockSpec
• Element with {padding=} is not supported.

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