jax-ml/jax · error · NotImplementedError

All block dimensions must be Elements or none of them can be

Error message

All block dimensions must be Elements or none of them can be Elements.

What it means

When any dimension of a block shape is a pallas_core.Element (an element-level windowed block), every other dimension must be an Element or Squeezed. Mixing Element dims with Blocked dims is unsupported because the TPU lowering can only emit an element_window attribute for uniformly element-wise blocks.

Source

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

          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
        ]
        pad_low, pad_high = map(list, zip(*padding))
        block_params["window_kind"] = ir.Attribute.parse(
            f"#tpu.element_window<{pad_low},{pad_high}>"
        )
      if pipeline_mode is not None:
        if not isinstance(pipeline_mode, pallas_core.Buffered):
          raise LoweringException(
              f"Unsupported pipeline mode: {pipeline_mode}."
          )
        if pipeline_mode.use_lookahead:
          raise NotImplementedError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make all dimensions uniformly Blocked (tile the whole shape) and handle padding inside the kernel body
If element semantics are needed, use Element/Squeezed for every dimension
Redesign the kernel: split the array so the element-windowed access is a separate rank-appropriate operand

Example fix

# before
block_shape=[Blocked(128), Element(padding=(1, 1))]

# after
block_shape=[Blocked(128), Blocked(1)]  # tile uniformly; pad in-kernel

Defensive patterns

Strategy: type-guard

Validate before calling

from jax._src.pallas import pallas_core
def check_uniform_element_blocks(bm):
    if any(isinstance(bd, pallas_core.Element) for bd in bm.block_shape):
        if not all(isinstance(bd, (pallas_core.Element, pallas_core.Squeezed)) for bd in bm.block_shape):
            raise ValueError('Element dims cannot be mixed with Blocked dims')

Type guard

def has_uniform_element_blocks(bm) -> bool:
    flags = [isinstance(bd, (pallas_core.Element, pallas_core.Squeezed))
             for bd in bm.block_shape]
    return all(flags) if any(isinstance(bd, pallas_core.Element) for bd in bm.block_shape) else True

Try / catch

try:
    pallas_call(...)
except NotImplementedError as e:
    if 'Elements or none' in str(e):
        switch all dims to Blocked and handle padding in-kernel

Prevention

Avoid Element descriptors unless the entire block is element-wise
Do padding/slicing inside the kernel body rather than via element windows

When it happens

Trigger: A block_shape like [Element(...), Blocked(16)] — some dims element-level, others tiled — reaching lower_jaxpr_into_pipelined_module; typically from custom BlockSpecs that use pallas_core.Element for one dimension only (e.g. to get per-element padding on the last dim while tiling the first).

Common situations: Hand-built block mappings using Element descriptors for fine-grained padding; experimental kernels combining element windows with tiling; internal code that wraps some dims in Element for dynamic shapes while leaving others Blocked.

Related errors

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