{"record":{"id":"8be66f1ede78c4c1","repo":"jax-ml/jax","slug":"reduce-axes-argument-to-vjp-is-deprecated","errorCode":null,"errorMessage":"reduce_axes argument to vjp is deprecated","messagePattern":"reduce_axes argument to vjp is deprecated","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1703,"sourceCode":"    ``vjpfun`` is a function from a cotangent vector with the same shape as\n    ``primals_out`` to a tuple of cotangent vectors with the same number and\n    shapes as ``primals``, representing the vector-Jacobian product of ``fun``\n    evaluated at ``primals``.\n\n  >>> import jax\n  >>>\n  >>> def f(x, y):\n  ...   return jax.numpy.sin(x), jax.numpy.cos(y)\n  ...\n  >>> primals, f_vjp = jax.vjp(f, 0.5, 1.0)\n  >>> xbar, ybar = f_vjp((-0.7, 0.3))\n  >>> print(xbar)\n  -0.61430776\n  >>> print(ybar)\n  -0.2524413\n  \"\"\"\n  if reduce_axes:\n    raise NotImplementedError(\"reduce_axes argument to vjp is deprecated\")\n  del reduce_axes\n  check_callable(fun)\n  canon = lambda x: x if isinstance(x, core.Tracer) else canonicalize_value(x)\n  primals_ft = ft.flatten(primals).map(canon)\n  primals_ft.map(dispatch.check_arg)\n  saveable = _saveable_args_flags(saveable_args, primals_ft.tree)\n  in_nzs_flat = None if in_nzs is None else tuptree_flags(\n      in_nzs, primals_ft.tree, 'in_nzs', 'the in_nzs argument to jax.vjp')\n  out_primals_ft, out_zeros, jaxpr, residuals, structured_residuals, *maybe_aux = \\\n      ad.linearize(fun, primals_ft, is_vjp=True, has_aux=has_aux,\n                   in_nzs=in_nzs_flat)\n\n  id_map = {id(x): i for i, x in enumerate(primals_ft)}\n  used, opaque_residuals = set(), []\n  spec = [used.add(id(r)) or RSpec(id_map[id(r)], True) if id(r) in id_map else\n          RSpec(opaque_residuals.append(r) or (len(opaque_residuals) - 1), False)\n          for r in residuals]\n  keep = lambda x, s: ((x if s else NotSaveable()) if id(x) in used else NotNeeded())","sourceCodeStart":1685,"sourceCodeEnd":1721,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1685-L1721","documentation":"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.","triggerScenarios":"Calling jax.vjp(f, x, reduce_axes=('batch',)) or any non-None/non-empty reduce_axes in current JAX.","commonSituations":"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.","solutions":["Remove the reduce_axes argument and perform reductions over named/mapped axes manually (e.g. with jax.lax.psum) inside or after the function","Migrate to modern SPMD (jit + NamedSharding) or vmap-based reductions","Pin an old JAX version only as a temporary compatibility stopgap"],"exampleFix":"# before\n_, vjp_fn = jax.vjp(loss_fn, x, reduce_axes=('batch',))\n# after\n_, vjp_fn = jax.vjp(psum_loss_fn, x)  # apply jax.lax.psum inside loss_fn","handlingStrategy":"validation","validationCode":"assert reduce_axes in (None, (), []), 'reduce_axes is removed from jax.vjp'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Grep codebase for reduce_axes before upgrading JAX","Move reductions into the function with jax.lax.psum","Track JAX changelogs for vjp API changes"],"tags":["jax","vjp","deprecated","migration"],"backgroundTag":"removed-api-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}