jax-ml/jax · error · ValueError

Block shape for {origin} (= {block_shape}) must have the sam

Error message

Block shape for {origin} (= {block_shape}) must have the same number of dimensions as the array shape {array_aval.shape}.

What it means

Each dimension of a BlockSpec's block_shape maps to one dimension of the corresponding array. to_block_mapping canonicalizes the block shape and requires len(block_shape) == len(array_aval.shape); otherwise there is no way to assign blocks to array dims and it raises this error naming the origin (input name).

Source

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

      if not hasattr(array_aval, "shape"):
        raise ValueError(
            "Array type must have a `shape` attribute, but got"
            f" {type(array_aval)}"
        )
    if self.index_map is None:
      index_map_func = default_index_map(len(array_aval.shape))
      index_map_dbg = api_util.debug_info("pallas_call index_map",
                                          default_index_map, (),{}
                                          )._replace(arg_names=("",) * len(index_map_avals))
      api_util.save_wrapped_fun_debug_info(index_map_func, index_map_dbg)
    else:
      index_map_func = self.index_map
    if self.block_shape is None:
      block_shape = _canonicalize_block_shape(array_aval.shape)
    else:
      block_shape = _canonicalize_block_shape(self.block_shape)
      if len(array_aval.shape) != len(block_shape):
        raise ValueError(
            f"Block shape for {origin} (= {block_shape}) "
            "must have the same number of dimensions as the "
            f"array shape {array_aval.shape}."
        )

    ref_block_shape = _get_ref_block_shape(block_shape)
    if isinstance(array_aval, jax_core.ShapedArray):
      arr_sh = array_aval.sharding
      ref_sharding = arr_sh.update(spec=arr_sh.spec.update(
          partitions=tuple(arr_sh.spec)[:len(ref_block_shape)]))
      block_array_aval = array_aval.update(
          shape=ref_block_shape, memory_space=jax_core.MemorySpace.Device,
          sharding=ref_sharding)
    elif isinstance(array_aval, state_types.AbstractLinVal):
      if not isinstance(array_aval.inner_aval, jax_core.ShapedArray):
        raise NotImplementedError  # TODO(mattjj,sharadmv)
      block_array_aval = array_aval.inner_aval.update(shape=ref_block_shape)
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make block_shape length equal the array ndim, using None for squeezed dims (e.g. (None, 128) for per-row blocks of a (B, 128) array)
2. Update index_map to return one index per block dim
3. Log array_aval.shape and block_shape side by side before pallas_call

Example fix

# before
spec = pl.BlockSpec(block_shape=(128,), index_map=lambda i: (i,))
# on array shape (4096, 128) -> error
# after
spec = pl.BlockSpec(block_shape=(None, 128), index_map=lambda i: (i,))

assert len(block_shape) == x.ndim, (
    f'block_shape {block_shape} rank != array ndim {x.ndim}')
spec = pl.BlockSpec(block_shape=block_shape, index_map=...)

def rank_matches(x, block_shape):
    return x.ndim == len(block_shape)

When it happens

Trigger: BlockSpec(block_shape=(128,), index_map=...) applied to a 2D array (e.g. shape (4096, 128)), or forgetting a trailing None for a squeezed dim (block_shape=(128, 128) on 3D data).

Common situations: Writing 1D-style kernels then applying them to 2D batches; adding a leading batch dim to inputs without updating the BlockSpec; copy-pasting specs between operands of different rank.

Related errors

• Stacking only supported when the block size along the stack
• Unsupported block dimension type: {type(dim)}. Allowed types
• Unsupported block shape type: {type(dim)}
• Array type must have a `shape` attribute, but got {type(arra
• Index map function {debug_info.func_src_info} for {origin} m

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