jax-ml/jax · error · NotImplementedError

Stack with Element indexing is not yet supported.

Error message

Stack with Element indexing is not yet supported.

What it means

Raised by _stack_pull_rule when propagating block specs through jnp.stack where the output block transform uses pallas_core.Element entries. Stack's block decomposition (block size along the new axis equals number of inputs) cannot be expressed with Element indexing.

Source

Thrown at jax/_src/pallas/fuser/block_spec.py:1866

      return util.tuple_update(idx, dimension, block_idx)

    return block_transform.replace(
        block_index_transform=new_block_index_transform
    )
  return [make_block_transform(i) for i in range(len(ctx.avals_in))]


@register_pull_block_spec_rule(lax.stack_p)
def _stack_pull_rule(
    ctx: PullRuleContext,
    block_transform: BlockIndexTransform,
    *,
    axis: int,
):
  block_shape = block_transform.block_shape
  is_element_block = [isinstance(bd, pallas_core.Element) for bd in block_shape]
  if any(is_element_block):
    raise NotImplementedError(
        'Stack with Element indexing is not yet supported.'
    )
  block_dim = block_shape[axis]
  if block_dim is None or isinstance(block_dim, pallas_core.Squeezed):
    block_dim = 1

  n = len(ctx.avals_in)
  if block_dim != n:
    raise NotImplementedError(
        "Stacking only supported when the block size along the stack axis "
        f"equals the number of inputs. Got block_dim={block_dim}, expected {n}."
    )

  new_block_shape = list(block_transform.block_shape)
  new_block_shape.pop(axis)

  def make_block_transform(child_index: int):
    def new_block_index_transform(*idxs):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Replace Element block entries with int/Squeezed/None descriptors in the output BlockSpec
Use concatenate or per-input writes instead of stack
Perform the stack outside the kernel and pass the result as a normally-blocked input

Example fix

# before
BlockSpec(block_shape=(pallas_core.Element(1), 64), ...)
out[...] = jnp.stack([a, b], axis=0)

# after
BlockSpec(block_shape=(2, 64), ...)
out[...] = jnp.stack([a, b], axis=0)

Defensive patterns

Strategy: validation

Validate before calling

import jax._src.pallas.pallas_core as pc
assert not any(isinstance(b, pc.Element) for b in block_shape), 'stack output cannot use Element blocks'

Type guard

def stack_safe_block_spec(block_shape):
    import jax._src.pallas.pallas_core as pc
    return not any(isinstance(b, pc.Element) for b in block_shape)

Prevention

Use int block sizes (equal to number of stacked inputs along the new axis) instead of Element
Prefer per-input writes or outside-kernel stacking

When it happens

Trigger: Using jnp.stack inside a pallas kernel where the stacked result feeds a ref whose BlockSpec uses pallas_core.Element along any dimension.

Common situations: Assembling per-block stacked outputs with fine-grained Element layouts; combining Mosaic Element indexing with stack-based output assembly.

Related errors

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