jax-ml/jax · error · TypeError

lax.associative_scan: fn argument should be callable.

Error message

lax.associative_scan: fn argument should be callable.

What it means

jax.lax.associative_scan requires its first argument to be a Python callable (the binary associative combine function). Before flattening the inputs, the implementation checks `callable(fn)` and raises TypeError immediately if it fails, since there is no meaningful way to perform a scan without a combine operation.

Source

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

  Example 2: partial products of an array of matrices

  >>> mats = jax.random.uniform(jax.random.key(0), (4, 2, 2))
  >>> partial_prods = lax.associative_scan(jnp.matmul, mats)
  >>> partial_prods.shape
  (4, 2, 2)

  Example 3: reversed partial sums of an array of numbers

  >>> lax.associative_scan(jnp.add, jnp.arange(0, 4), reverse=True)
  Array([6, 6, 5, 3], dtype=int32)

  .. [BLE1990] Blelloch, Guy E. 1990. "Prefix Sums and Their Applications.",
    Technical Report CMU-CS-90-190, School of Computer Science, Carnegie Mellon
    University.
  """
  if not callable(fn):
    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 "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a callable as the first argument, e.g. jnp.add or a lambda: lax.associative_scan(jnp.add, elems)
2. If you wanted a string op, wrap it: lax.associative_scan(lambda a, b: a + b, elems)
3. Check argument order — signature is associative_scan(fn, elems, reverse=False)

Example fix

// before
lax.associative_scan('+', elems)
// after
lax.associative_scan(jnp.add, elems)

assert callable(fn), 'fn must be callable'

def is_scan_fn(fn) -> bool:
    return callable(fn)

When it happens

Trigger: Calling lax.associative_scan(fn, elems) where fn is not callable, e.g. passing a string like '+', an operator symbol, a jnp array, or forgetting the function argument entirely and passing only arrays.

Common situations: Developers coming from other libraries where reductions accept string op names (e.g. numpy/jnp reductions or torch); passing jnp.add result instead of jnp.add; refactoring code and dropping the fn argument.

Related errors

• shard_map requires a callable for its first argument, but go
• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.

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