jax-ml/jax · error · NotImplementedError

Please file an issue at https://github.com/jax-ml/jax/issues

Error message

Please file an issue at https://github.com/jax-ml/jax/issues

What it means

`all_gather_reduced` has no batching (vmap) rule registered, so applying `jax.vmap` (or any API that uses fancy batching, like scan with batched carry) over a function containing it raises NotImplementedError with a request to file an issue. The batching rule `_all_gather_reduced_batched_collective` is a stub that always raises. It is a known API gap, not a user-logic error.

Source

Thrown at jax/_src/lax/parallel.py:2647

      platform=platform, is_async=is_async)

mlir.register_lowering(all_gather_reduced_p, _all_gather_reduced_lowering)
for p in ("cuda", "rocm", "tpu"):
  mlir.register_lowering(all_gather_reduced_p,
                         partial(_all_gather_reduced_lowering, platform=p),
                         platform=p)

def _all_gather_reduced_transpose_rule(
    cts, x, *, all_gather_dimension, axis_name, axis_size, tiled):
  return (unreduced_psum_scatter(cts, axis_name=axis_name,
                                 scatter_dimension=all_gather_dimension,
                                 tiled=tiled),)
ad.deflinear2(all_gather_reduced_p, _all_gather_reduced_transpose_rule)

def _all_gather_reduced_batched_collective(
    axis_data, vals_in, dims_in, all_gather_dimension, axis_name, axis_size,
    tiled):
  raise NotImplementedError(
      "Please file an issue at https://github.com/jax-ml/jax/issues")
batching.fancy_primitive_batchers[all_gather_reduced_p] = _all_gather_reduced_batched_collective

####################### unreduced_psum_scatter ###########################

# Unreduced -> Varying collective
def unreduced_psum_scatter(x, axis_name, *, scatter_dimension=0, tiled=False):
  if not isinstance(axis_name, tuple):
    axis_name = (axis_name,)
  if not axis_name:
    return x
  axis_size = _axis_size(axis_name, None)
  def bind(leaf):
    return unreduced_reduce_scatter_p.bind(
        leaf, axis_name=axis_name, scatter_dimension=scatter_dimension,
        axis_size=axis_size, tiled=tiled)
  return tree_util.tree_map(bind, x)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Move the collective outside the vmap'd function (manual batching via a leading axis)
Use in_axes=None for the collective arguments or restructure with explicit loop/scan semantics
File an issue at https://github.com/jax-ml/jax/issues as the message requests and pin a JAX version behavior
Check newer JAX versions where a batching rule may have been added

Example fix

// before
f = jax.vmap(lambda x: lax.all_gather_reduced(x, axis_name='i', all_gather_dimension=0, axis_size=8))
y = f(xs)
// after
def body(x):
    return lax.all_gather_reduced(x, axis_name='i', all_gather_dimension=0, axis_size=8)
y = jax.vmap(body, in_axes=(None,), out_axes=None)(xs)  # or hoist collective out of vmap

Defensive patterns

Strategy: fallback

Validate before calling

null

Type guard

null

Try / catch

try:
    ys = jax.vmap(f)(xs)
except NotImplementedError as e:
    if 'all_gather_reduced' in str(e):
        ys = jax.lax.map(f, xs)  # sequential fallback
    else:
        raise

Prevention

Hoist collectives out of vmap'd functions
Check for a registered batching rule before vmap-ing exotic primitives
Watch jax release notes for all_gather_reduced batching support

When it happens

Trigger: Wrapping a function that calls `lax.all_gather_reduced` in `jax.vmap`; batched autodiff or `lax.scan` that triggers the fancy batching machinery on this primitive.

Common situations: Ensembling / batched training loops where a per-example function contains a collective; using vmap to auto-batch code written for the named-axis mesh API.

Related errors

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