jax-ml/jax · error · NotImplementedError

reduce_axes argument to vjp is deprecated

Error message

reduce_axes argument to vjp is deprecated

What it means

jax.vjp once accepted a reduce_axes keyword for collecting vjp contributions across named axes; it has been removed and now raises NotImplementedError immediately if a truthy value is passed.

Source

Thrown at jax/_src/api.py:1703

    ``vjpfun`` is a function from a cotangent vector with the same shape as
    ``primals_out`` to a tuple of cotangent vectors with the same number and
    shapes as ``primals``, representing the vector-Jacobian product of ``fun``
    evaluated at ``primals``.

  >>> import jax
  >>>
  >>> def f(x, y):
  ...   return jax.numpy.sin(x), jax.numpy.cos(y)
  ...
  >>> primals, f_vjp = jax.vjp(f, 0.5, 1.0)
  >>> xbar, ybar = f_vjp((-0.7, 0.3))
  >>> print(xbar)
  -0.61430776
  >>> print(ybar)
  -0.2524413
  """
  if reduce_axes:
    raise NotImplementedError("reduce_axes argument to vjp is deprecated")
  del reduce_axes
  check_callable(fun)
  canon = lambda x: x if isinstance(x, core.Tracer) else canonicalize_value(x)
  primals_ft = ft.flatten(primals).map(canon)
  primals_ft.map(dispatch.check_arg)
  saveable = _saveable_args_flags(saveable_args, primals_ft.tree)
  in_nzs_flat = None if in_nzs is None else tuptree_flags(
      in_nzs, primals_ft.tree, 'in_nzs', 'the in_nzs argument to jax.vjp')
  out_primals_ft, out_zeros, jaxpr, residuals, structured_residuals, *maybe_aux = \
      ad.linearize(fun, primals_ft, is_vjp=True, has_aux=has_aux,
                   in_nzs=in_nzs_flat)

  id_map = {id(x): i for i, x in enumerate(primals_ft)}
  used, opaque_residuals = set(), []
  spec = [used.add(id(r)) or RSpec(id_map[id(r)], True) if id(r) in id_map else
          RSpec(opaque_residuals.append(r) or (len(opaque_residuals) - 1), False)
          for r in residuals]
  keep = lambda x, s: ((x if s else NotSaveable()) if id(x) in used else NotNeeded())

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the reduce_axes argument and perform reductions over named/mapped axes manually (e.g. with jax.lax.psum) inside or after the function
2. Migrate to modern SPMD (jit + NamedSharding) or vmap-based reductions
3. Pin an old JAX version only as a temporary compatibility stopgap

Example fix

# before
_, vjp_fn = jax.vjp(loss_fn, x, reduce_axes=('batch',))
# after
_, vjp_fn = jax.vjp(psum_loss_fn, x)  # apply jax.lax.psum inside loss_fn

assert reduce_axes in (None, (), []), 'reduce_axes is removed from jax.vjp'

When it happens

Trigger: Calling jax.vjp(f, x, reduce_axes=('batch',)) or any non-None/non-empty reduce_axes in current JAX.

Common situations: Running old code or tutorials written against JAX from ~2020 that used reduce_axes for pmap-compatible vjps; copy-pasted training loops with named-axis vjp logic.

Related errors

• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• But the tree structures differ:
• the VJP function was applied before restoring its not-saveab

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