jax-ml/jax · error · NotImplementedError

Unsupported block size type: {type(bd)}

Error message

Unsupported block size type: {type(bd)}

What it means

When computing block sizes in the Mosaic GPU pipeline, a block dimension must be an int, pl.Blocked, or pl.Element. Any other type (e.g. None, custom object, or a pl-pipeline placeholder) raises NotImplementedError with the offending type name.

Source

Thrown at jax/_src/pallas/mosaic_gpu/pipeline.py:63

OOBFillMode = gpu_primitives.OOBFillMode

class PipelineCallback[T](Protocol):
  """A callback that returns the same type as the input."""
  def __call__(self, arg: T, /) -> T: ...

type BlockSpecPytree = Sequence[pl.BlockSpec | BlockSpecPytree]
type AbstractRefPytree = Sequence[state.AbstractRef | AbstractRefPytree]


def _get_block_size(bd: pl.BlockDim | int | None) -> int:
  match bd:
    case int():
      return bd
    case pl.Blocked() | pl.Element():
      return bd.block_size
    case _:
      raise NotImplementedError(f"Unsupported block size type: {type(bd)}")


def _get_block_shape(spec: pallas_core.BlockSpec, ref_shape: tuple[int, ...]):
  if spec.block_shape is None:
    return ref_shape

  block_shape = tuple(
      _get_block_size(bd)
      for bd in spec.block_shape
      if not (bd is None or isinstance(bd, pl.Squeezed))
  )
  return block_shape


def _is_fully_in_bounds(
    spec: pallas_core.BlockSpec, operand_shape: tuple[int, ...]
) -> bool:
  """Returns whether all windowed accesses into an operand stay in bounds."""

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use concrete ints, pl.Blocked(...), or pl.Element(...) in block_shape
2. Ensure grid/blockspec construction goes through public pallas APIs so placeholders are resolved
3. Upgrade JAX if the value comes from internal pipeline plumbing

Example fix

# before
bs = pl.BlockSpec(block_shape=(None, 128))
# after
bs = pl.BlockSpec(block_shape=(ref_len, 128))  # concrete ints / pl.Blocked

from jax.experimental import pallas as pl
assert all(isinstance(b, (int, pl.Blocked, pl.Element)) for b in block_shape)

def valid_block_spec(bd) -> bool:
    from jax.experimental import pallas as pl
    return isinstance(bd, (int, pl.Blocked, pl.Element))

When it happens

Trigger: Constructing a BlockSpec whose block_shape entries are unusual types — commonly a None or an unspecialized placeholder from using pallas pipeline APIs before grid specialization resolves them.

Common situations: Passing None block_shapes expecting full-shape blocks in contexts that call _get_block_size before defaults are applied; using stale pl.BlockSpec argument shapes across JAX version changes.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Sparse metadata format not implemented for {operand_dtype=}

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