jax-ml/jax · error · NotImplementedError

Unsupported block dimension type: {type(bd)} for block shape

Error message

Unsupported block dimension type: {type(bd)} for block shape: {bm.block_shape}

What it means

Block shapes in Pallas grid mappings must be built from pallas_core.Element, pallas_core.Squeezed, or pallas_core.Blocked dimension descriptors. If any dimension of bm.block_shape is another type (e.g. a plain int from manual BlockSpec construction), the TPU lowering raises NotImplementedError listing the offending type.

Source

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

          tpu_memory_space == tpu_core.MemorySpace.VMEM
          and bm.has_trivial_window()
      ):
        pipeline_mode = pallas_core.Buffered(1)

      # If we have an extended dtype, we need to add the block shape for the
      # remaining physical dtype.
      block_shape += list(_get_aval_physical_dtype_shape(bm.block_aval.inner_aval))
      block_shape = dynamic_shape_replacement_fn(block_shape)
      window_shape = ir.DenseI64ArrayAttr.get(block_shape)
      block_params: dict[str, ir.Attribute] = dict(
          window_bounds=window_shape,
          transform_indices=ir.FlatSymbolRefAttr.get(func_name),
      )
      for bd in bm.block_shape:
        if not isinstance(
            bd, (pallas_core.Element, pallas_core.Squeezed, pallas_core.Blocked)
        ):
          raise NotImplementedError(
              "Unsupported block dimension type: "
              f"{type(bd)} for block shape: {bm.block_shape}"
          )
      is_element_block = [isinstance(bd, pallas_core.Element)
                          for bd in bm.block_shape]
      if any(is_element_block):
        is_element_or_squeezed_block = [
            isinstance(bd, (pallas_core.Element, pallas_core.Squeezed))
            for bd in bm.block_shape
        ]
        if not all(is_element_or_squeezed_block):
          raise NotImplementedError(
              "All block dimensions must be Elements or none of them can be"
              " Elements."
          )
        padding = [
            bd.padding if isinstance(bd, pallas_core.Element) else (0, 0)
            for bd in bm.block_shape

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Build block shapes with the public API (BlockSpec with plain tuples) so JAX constructs Blocked/Element descriptors for you
If manipulating internals, convert each dim via pallas_core.Blocked(...) or keep the original descriptor objects
Update custom code that assumed integer block dims after a JAX upgrade; check the BlockSpec constructor signature in your JAX version
Avoid hand-assembling GridMapping; use pallas.create_grid / pallas_call's automatic mapping instead

Example fix

# before
GridMapping(block_mappings=[BlockMapping(block_shape=[16, 16], ...)])

# after
# let pallas_call build mappings from BlockSpec:
pallas_call(kernel, BlockSpec(block_shape=(16, 16), index_map=...), grid=(...))

Defensive patterns

Strategy: type-guard

Validate before calling

from jax._src.pallas import pallas_core
def check_block_shape_types(bm):
    for bd in bm.block_shape:
        if not isinstance(bd, (pallas_core.Element, pallas_core.Squeezed, pallas_core.Blocked)):
            raise TypeError(f'bad block dim {bd!r} in {bm.block_shape}')

Type guard

def is_valid_block_shape(bm) -> bool:
    return all(isinstance(bd, (pallas_core.Element, pallas_core.Squeezed, pallas_core.Blocked))
               for bd in bm.block_shape)

Try / catch

try:
    lower_jaxpr_into_pipelined_module(...)
except NotImplementedError as e:
    if 'Unsupported block dimension type' in str(e):
        rebuild mappings via public BlockSpec and retry

Prevention

Never hand-construct GridMapping/BlockMapping; use BlockSpec + pallas_call
After JAX upgrades, re-verify internal APIs your code touches
Keep block shapes as plain tuples at the user level

When it happens

Trigger: Constructing BlockSpec/BlockMapping objects manually (or via internal APIs) with block_shape given as raw ints/strings instead of pallas_core.Blocked(...) / Element(...); intercepting or transforming block mappings with custom code that replaces descriptors with ints.

Common situations: Writing custom interpreters/partials over Pallas internals; monkey-patching or serializing/deserializing block mappings; upgrading JAX versions where block_shape representation changed from ints to descriptor objects, breaking custom glue code.

Related errors

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