jax-ml/jax · error · ValueError

Element with {padding=} is not supported.

Error message

Element with {padding=} is not supported.

What it means

In JAX's Mosaic GPU pipeline (jax.experimental.pallas / tpu), each dimension of a BufferedRef's windowing spec is described by a BlockDimension. When a dimension uses Element(block_size, padding=...) with non-zero padding, compute_slice cannot translate it into a dynamic_slice and raises this ValueError. Zero padding is fine; any other padding tuple is rejected.

Source

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

  def unbind_refs(self):
    if not self.is_buffered and not self.has_allocated_buffer:
      return dataclasses.replace(self, window_ref=None)
    return self

  def compute_slice(self, grid_indices):
    """Compute the indexers for the window at given grid indices."""
    indices = self.compute_index(*grid_indices)
    assert self.block_shape is not None
    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."
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change padding=(left, right) on Element block dimensions to (0, 0) and handle out-of-bounds edges by adjusting the index map / block shape
2. Use a BoundedSlice or explicit start indices to express the padded window instead of Element padding
3. Run in interpreter mode (jax.platform or interpret=True) if padding semantics are required and pipeline support isn't needed

Example fix

# before
BlockSpec(index_map=lambda i: (i * 128,), block_shape=(128,), padding=((2, 2),))
# after
BlockSpec(index_map=lambda i: (i * 128 - 2,), block_shape=(132,), padding=((0, 0),))

bad = [d for d in spec.block_shape if isinstance(d, Element) and d.padding != (0, 0)]
assert not bad, f'Unsupported Element padding: {bad}'

from jax._src.pallas.mosaic.pipeline import Element
def has_zero_padding(spec) -> bool:
    return all(not isinstance(d, Element) or d.padding == (0, 0)
               for d in spec.block_shape)

When it happens

Trigger: Passing a BlockSpec whose block dimensions include an Element with padding != (0, 0) to a kernel using pipelined/buffered execution (e.g. emit_pipeline with buffering, fetch_with_lookahead, or make_output_bref paths that build indexers via compute_slice).

Common situations: Converting a Pallas kernel from CPU/TPU interp mode to GPU Mosaic pipeline mode where padded BlockSpecs were previously tolerated; using halos/borders expressed as Element padding instead of explicit slice bounds.

Related errors

• Padding with aliasing not supported.
• Acc ref must be at least 2D, got shape {shape}
• 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/aa660c13b361a9f6. Report an issue: GitHub.