jax-ml/jax · error · ValueError

Unsupported block shape type

Error message

Unsupported block shape type: {type(dim)}

What it means

Internal helper _get_block_dim_size extracts the size from a BlockDim via match: Squeezed/None -> 0, Blocked/Element/BoundedSlice/Indirect -> their block_size, plain int -> itself. Any other type falls to the wildcard case and raises this ValueError.

Solutions

1. Ensure every entry of block_shape is int, None, Blocked, Squeezed, Element, BoundedSlice, or Indirect from the same JAX version
2. Don't create custom block-dim classes; express custom tiling with index_map instead
3. Reinstall/align JAX versions if the values come from library internals

from jax._src.pallas import core as pcore
allowed = (int, type(None), pcore.Blocked, pcore.Squeezed, pcore.Element,
           pcore.BoundedSlice, pcore.Indirect)
assert all(isinstance(d, allowed) for d in block_shape)

def is_valid_block_shape(shape):
    return all(d is None or isinstance(d, (int, Blocked, Squeezed, Element, BoundedSlice, Indirect)) for d in shape)

When it happens

Trigger: Passing a custom or malformed object inside a canonicalized block_shape — typically from constructing BlockSpec with non-standard dims or from internal code paths (_get_block_shape, has_trivial_window, triton lowering) receiving corrupted block specs.

Common situations: Subclassing or monkey-patching BlockDim types; JAX internal version mismatches; feeding block shapes built by older/newer pallas APIs.

Related errors

• Array type must have a `shape` attribute, but got
• Block shape for (= ) must have the same number of…
• BlockMapping for has captured constants
• Cannot pull iota along dimension
• Index map function for must return values to match …

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

Appendix: source

Thrown at jax/_src/pallas/core.py:486

def _canonicalize_block_shape(block_shape: Sequence[BlockDim | int | None]
                              ) -> tuple[BlockDim, ...]:
  return tuple(_canonicalize_block_dim(dim) for dim in block_shape)


def _get_block_dim_size(dim: BlockDim) -> int:
  match dim:
    case Squeezed():
      return 1
    case (
        Blocked(block_size)
        | Element(block_size)
        | BoundedSlice(block_size)
        | Indirect(block_size)
    ):
      return block_size
    case _:
      raise ValueError(f"Unsupported block shape type: {type(dim)}")

def get_block_size(dim: BlockDim | int | None) -> int:
  match dim:
    case int():
      return dim
    case Squeezed() | None:
      return 1
    case (
        Blocked(block_size)
        | Element(block_size)
        | BoundedSlice(block_size)
        | Indirect(block_size)
    ):
      return block_size
    case _:
      raise ValueError(f"Unsupported block shape type: {type(dim)}")

View on GitHub (pinned to 1e1c6a8fc0)