jax-ml/jax · error · ValueError

Cannot pull iota along dimension {dimension} with None block

Error message

Cannot pull iota along dimension {dimension} with None block size.

What it means

A ValueError (not NotImplementedError) from the iota-pulling logic: when the fuser tries to materialize an iota (index/identifier array) into a kernel argument along a dimension whose block size is None (unbounded), it cannot compute per-block index offsets, so it refuses.

Source

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

  local_iota = jax.lax.broadcasted_iota(dtype, iota_shape, dim_)
  return local_iota + block_idx[dimension] * _block_size(
      block_spec.block_shape[dimension]
  )


@register_pull_block_spec_rule(lax.iota_p)
def _iota_pull_rule(
    ctx: PullRuleContext,
    block_transform: BlockIndexTransform,
    *,
    dtype: jnp.dtype,
    dimension: int,
    shape: tuple[int, ...],
    sharding: jax.sharding.Sharding,
):
  del ctx, sharding, dtype, shape
  if block_transform.block_shape[dimension] is None:
    raise ValueError(
        f'Cannot pull iota along dimension {dimension} with None block size.'
    )
  return []


def _pattern_match_lanes_to_sublanes_reshape(
    aval_in: core.ShapedArray,
    aval_out: core.ShapedArray,
) -> bool:
  # Pattern matches a reshape of the form (..., n * l) -> (..., n, l)
  # where l is a multiple of 128.

  *leading_out, last_dim_in = aval_in.shape
  *leading_in, second_to_last_dim_out, last_dim = aval_out.shape
  if leading_in != leading_out:
    return False
  if second_to_last_dim_out * last_dim != last_dim_in:
    return False

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Set a concrete integer block size on the iota's dimension in the BlockSpec
Generate the iota/arange outside the kernel and pass it as a regular block-mapped input
Derive in-kernel indices from index_map arguments (program ids) instead of a materialized iota

Example fix

// before
spec = BlockSpec((None, d), ...)  # iota pulled along axis 0 which is None
pos = jnp.arange(seq_len)  # becomes iota
// after
spec = BlockSpec((128, d), ...)  # concrete block size on axis 0

Defensive patterns

Strategy: validation

Validate before calling

assert spec.block_shape[iota_axis] is not None, 'iota axis needs a concrete block size'

Type guard

def iota_axis_ok(block_shape, dimension) -> bool:
    return block_shape[dimension] is not None

Try / catch

try:
    out = fused_fn(x)
except ValueError as e:
    if 'Cannot pull iota' in str(e):
        out = fused_fn(x, idx=jnp.arange(n))  # pass iota as explicit input
    else:
        raise

Prevention

Materialize arange/iota outside kernels with unbounded axes
Derive indices from index_map program ids
Avoid None block dims on axes used for positional data

When it happens

Trigger: Code paths that pull lax.iota / index materialization (e.g. from jnp.arange-like patterns or positional encodings) into a fused Pallas operand where block_transform.block_shape[dimension] is None along the iota's axis.

Common situations: Using arange/index-derived values (RoPE positions, masks) with BlockSpecs that leave the sequence axis unbounded; dynamic-shape kernels where one axis is None for flexibility; grid computed from data rather than block shapes.

Related errors

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