jax-ml/jax · error · ValueError

Unsupported block dimension type: {type(dim)}. Allowed types

Error message

Unsupported block dimension type: {type(dim)}. Allowed types: `pl.Squeezed`, `pl.Blocked`, `pl.Element`, `int`, `None`.

What it means

Block shapes in BlockSpec must consist of ints, None, or the symbolic dims pl.Blocked/pl.Squeezed/pl.Element. _canonicalize_block_dim tries each form and, when the value can't be interpreted (failed int() conversion), raises this ValueError listing the allowed types.

Source

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


def _canonicalize_block_dim(dim: BlockDim | int | None) -> BlockDim:
  match dim:
    case None:
      return squeezed
    case int():
      return Blocked(int(dim))
    case Squeezed() | Blocked() | Element() | BoundedSlice() | Indirect():
      return dim
    case _:
      # Handle case where the dim is a symbolic dimension so we assume it is
      # Blocked.
      if jax_core.is_symbolic_dim(dim):
        return Blocked(dim)
      try:
        return Blocked(int(dim))
      except Exception as e:
        raise ValueError(
            f"Unsupported block dimension type: {type(dim)}. Allowed types:"
            " `pl.Squeezed`, `pl.Blocked`, `pl.Element`, `int`, `None`."
        ) from e

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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use integer division or int() when computing block dims: n // 8 or cdiv helper
Keep block_shape entries as Python ints, None, or pl.Blocked/Squeezed/Element only
Print block_shape right before BlockSpec construction to spot float/str entries

Example fix

# before
block = (x.shape[0] / tile, None)  # floats
spec = pl.BlockSpec(block_shape=block, index_map=...)
# after
block = (x.shape[0] // tile, None)
spec = pl.BlockSpec(block_shape=block, index_map=...)

Defensive patterns

Strategy: validation

Validate before calling

dims = tuple(int(d) if isinstance(d, (int, float)) and not isinstance(d, bool) else d for d in block_shape)
block_shape = tuple(d if isinstance(d, (int, type(None))) or hasattr(d, 'block_size') or hasattr(d, 'size') else int(d) for d in dims)

Type guard

def valid_block_dim(d):
    import jax.experimental.pallas as pl
    return d is None or isinstance(d, int) or isinstance(d, (pl.Blocked, pl.Squeezed, pl.Element))

Prevention

Use // not / when computing tile sizes
Run a small assertion pass over block_shape before BlockSpec

When it happens

Trigger: Passing a BlockSpec block_shape containing a float, numpy scalar that fails int conversion, string, or arbitrary object — e.g. block_shape=(102.0, None) instead of (1024, None), or a JAX tracer inside block_shape.

Common situations: Computing block sizes with division that yields floats (n / 8 instead of n // 8); passing weakly-typed numpy floats; typos in BlockSpec tuples.

Related errors

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