jax-ml/jax · error · ValueError

Array inputs to associative_scan must have the same first di

Error message

Array inputs to associative_scan must have the same first dimension. (saw: {})

What it means

associative_scan flattens the (possibly pytree) input and requires every leaf array to have the same length along the scan axis so the combine tree is well-defined. After computing num_elems from the first leaf, it verifies all others match and raises ValueError listing all leaf shapes otherwise.

Source

Thrown at jax/_src/lax/control_flow/loops.py:2900

  def combine(a_flat, b_flat):
    # Lower `fn` to operate on flattened sequences of elems.
    a = tree_unflatten(tree, a_flat)
    b = tree_unflatten(tree, b_flat)
    c = fn(a, b)
    c_flat, _ = tree_flatten(c)
    return c_flat

  # Check that all inputs have a consistent leading dimension `num_elems`.
  axis = util.canonicalize_axis(axis, elems_flat[0].ndim)

  if not core.is_constant_dim(elems_flat[0].shape[axis]):
    raise NotImplementedError("associative scan over axis "
        f"of non-constant size: {elems_flat[0].shape[axis]}. You may be "
        "able to avoid this on TPU. See b/274176030.")
  num_elems = int(elems_flat[0].shape[axis])
  if not all(int(elem.shape[axis]) == num_elems for elem in elems_flat[1:]):
    raise ValueError('Array inputs to associative_scan must have the same '
                     'first dimension. (saw: {})'
                     .format([elem.shape for elem in elems_flat]))


  # Summary of algorithm:
  #
  # Consider elements of `_scan(elems)` at odd indices. That's the same as first
  # summing successive pairs of elements of `elems` and performing a scan on
  # that half sized tensor. We perform the latter scan by recursion.
  #
  # Now consider the even elements of `_scan(elems)`. These can be computed
  # from the odd elements of `_scan(elems)` by adding each odd element of
  # `_scan(elems)` to the matching even element in the original `elems`.
  #
  # We return the odd and even elements interleaved.
  #
  # For the base case of the recursion we return the first element
  # of `elems` followed by the sum of the first two elements computed as

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure every array in the pytree has identical size along the scan axis (pad or slice all members consistently)
2. If elements are independent, run separate associative_scan calls per array
3. Print [a.shape for a in jax.tree_util.tree_leaves(elems)] to find the mismatched leaf

Example fix

// before
lax.associative_scan(fn, (jnp.ones(5), jnp.ones(4)))
// after
lax.associative_scan(fn, (jnp.ones(5), jnp.ones(5)))

import jax.tree_util as jtu
leaves = jtu.tree_leaves(elems)
n = leaves[0].shape[axis]
assert all(l.shape[axis] == n for l in leaves), [l.shape for l in leaves]

When it happens

Trigger: Passing a pytree of arrays (tuple/dict) whose members have different lengths along `axis`, e.g. ((jnp.ones(5), jnp.ones(3)),) — note only the scanned axis must match, other dims may differ; or mixing arrays of different sequence lengths in a structured carry.

Common situations: Scanning over structured state (e.g. (cumsum, logprobs) tuples) where one element was sliced or padded differently; off-by-one slicing like x[:-1] applied to only one member of the tuple.

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

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