jax-ml/jax · error · NotImplementedError

associative scan over axis of non-constant size: {}. You may

Error message

associative scan over axis of non-constant size: {}. You may be able to avoid this on TPU. See b/274176030.

What it means

associative_scan implements a work-efficient parallel prefix scan whose lowering depends on knowing the scan-axis length at compile (trace) time. If the scanned axis has a non-constant (polynomial/dynamic) size, e.g. under jax.export or with shape-poly dimensions, JAX cannot build the fixed combine tree and raises NotImplementedError, pointing to b/274176030 for TPU workarounds.

Source

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

    raise TypeError("lax.associative_scan: fn argument should be callable.")
  elems_flat, tree = tree_flatten(elems)

  if reverse:
    elems_flat = [lax.rev(elem, [axis]) for elem in elems_flat]

  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`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Make the scanned axis a static dimension (exclude it from polymorphic dims) so its size is known at trace time
  2. Pad/truncate the sequence to a fixed length before the scan
  3. Rewrite as an explicit lax.fori_loop / lax.scan over the dynamic axis, which tolerates non-constant bounds on TPU
  4. Follow b/274176030 for the TPU-specific workaround

Example fix

// before
exp.export(fn, polynomial_shapes=('[n,]',))  # n dynamic, then associative_scan over axis 0
// after
exp.export(fn, polynomial_shapes=('(_,_)',))  # make scan axis static, pad input to fixed length
Defensive patterns

Strategy: validation

Validate before calling

import jax
size = jax.api shapes... # check statically:
# x = jnp.ones((n, d))
assert jnp.asarray(x).shape[axis] is not a dynamic dim — simplest: ensure shape ints known:
assert isinstance(x.shape[axis], int)

Try / catch

try:
    lax.associative_scan(fn, elems)
except NotImplementedError as e:
    if 'non-constant size' in str(e):
        out = lax.fori_loop_based_scan(fn, elems)  # fallback explicit loop
    else:
        raise

Prevention

When it happens

Trigger: Scanning along an axis whose dimension is a symbolic/polynomial dimension (jax.experimental.export, shape polymorphism) or otherwise not statically known, e.g. associative_scan(fn, x) where x.shape[axis] is a DimExpr rather than an int.

Common situations: Using jax.export with polymorphic shapes and then calling the exported artifact with dynamic sequence lengths; migrating scan code to JAX shape-polymorphic pipelines; running on TPU where dynamic-size loops are otherwise supported.

Related errors

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