jax-ml/jax · error · ValueError

Unsupported block dimension type: {type(bd)}

Error message

Unsupported block dimension type: {type(bd)}

What it means

computeSlice converts each BlockDimension of a BlockSpec into a dynamic_slice indexer via Python match. Only None, Squeezed, Element, BoundedSlice, Blocked, and int are supported; any other type falls through to this ValueError.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:838

    assert len(self.block_shape) == len(indices)
    indexer = []
    for bd, idx in zip(self.block_shape, indices, strict=True):
      match bd:
        case None | Squeezed():
          # Dimension is squeezed out so we don't do anything.
          indexer.append(idx)
        case Element(block_size, padding=padding):
          if padding != (0, 0):
            raise ValueError(f"Element with {padding=} is not supported.")
          indexer.append(ds(idx, block_size))
        case BoundedSlice(block_size):
          indexer.append(ds(idx.start, block_size))
        case Blocked(block_size):
          indexer.append(ds(idx * block_size, block_size))
        case int():
          indexer.append(ds(idx * bd, bd))
        case _:
          raise ValueError(f"Unsupported block dimension type: {type(bd)}")
    return tuple(indexer)

  def initialize_slots(self) -> BufferedRef:
    if self.window_ref is None and self.prefetched_count > 0:
      raise ValueError(
          "Expected external window buffer to be bound for prefetched input "
          f"(prefetched_count={self.prefetched_count}), but window_ref is None. "
          "Ensure .with_window_ref(...) is called on the BufferedRef in allocations."
      )
    return dataclasses.replace(
        self,
        copy_in_slot=jnp.uint32(0) if self.buffer_type.is_input else None,
        wait_in_slot=jnp.uint32(0) if self.buffer_type.is_input else None,
        copy_out_slot=jnp.uint32(0) if self.buffer_type.is_output else None,
        wait_out_slot=jnp.uint32(0) if self.buffer_type.is_output else None,
        next_fetch=(
            tuple(jnp.int32(0) for _ in range(self._grid_rank))
            if self._grid_rank is not None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure every block dimension is a Python int or an instance of a supported type (Blocked, BoundedSlice, Element, Squeezed, None)
2. Convert numpy scalars with int(...) before constructing the BlockSpec
3. If you subclassed a block dimension type, replace it with the closest supported type

Example fix

# before
block_shape=(np.int64(128),)
# after
block_shape=(int(np.int64(128)),)

import numpy as np
block_shape = tuple(int(d) if isinstance(d, np.integer) else d for d in block_shape)

SUPPORTED = (int, Element, BoundedSlice, Blocked, Squeezed, type(None))
def dims_supported(block_shape) -> bool:
    return all(d is None or isinstance(d, SUPPORTED) and not isinstance(d, bool) or d is None for d in block_shape)

When it happens

Trigger: Passing a custom or invalid object as a block dimension in a BlockSpec, e.g. a numpy integer (np.int64) instead of a Python int, or a future/unsupported BlockDimension subclass, into a pipelined Mosaic kernel.

Common situations: Using np.int64/np.intp values (from numpy arrays or computed sizes) as block_shape entries; defining a custom BlockDimension subclass expecting it to be honored.

Related errors

• MemoryRef type must be a ShapedArray, got {type(self.inner_a
• Unsupported block dimension type: {type(dim)}. Allowed types
• Out-of-bounds block index {block_indices} for input "{input_
• Out-of-bounds block index {block_indices} for output "{outpu
• Unsupported block dim type: {type(b)}

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