{"record":{"id":"cc4dae83b47e554b","repo":"jax-ml/jax","slug":"reduce-axes-argument-to-transpose-is-deprecated","errorCode":null,"errorMessage":"reduce_axes argument to transpose is deprecated","messagePattern":"reduce_axes argument to transpose is deprecated","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":2094,"sourceCode":"      See below for an example. (Note that the duck-typed objects cannot be\n      namedtuples because those are treated as standard Python containers.)\n\n  Returns:\n    A callable that calculates the transpose of ``fun``. Valid input into this\n    function must have the same shape/dtypes/structure as the result of\n    ``fun(*primals)``. Output will be a tuple, with the same\n    shape/dtypes/structure as ``primals``.\n\n  >>> import jax\n  >>>\n  >>> f = lambda x, y: 0.5 * x - 0.5 * y\n  >>> scalar = jax.ShapeDtypeStruct(shape=(), dtype=np.dtype(np.float32))\n  >>> f_transpose = jax.linear_transpose(f, scalar, scalar)\n  >>> f_transpose(1.0)\n  (Array(0.5, dtype=float32), Array(-0.5, dtype=float32))\n  \"\"\"\n  if reduce_axes:\n    raise NotImplementedError(\"reduce_axes argument to transpose is deprecated\")\n  del reduce_axes\n  primals_flat, in_tree = tree_flatten(primals)\n  flat_fun, out_tree = flatten_fun_nokwargs(\n      lu.wrap_init(fun,\n                   debug_info=debug_info(\"linear_transpose\", fun, primals, {})),\n      in_tree)\n  in_avals = [shaped_abstractify(x) for x in primals_flat]\n  in_dtypes = [a.dtype for a in in_avals if not a.is_high]\n\n  in_pvals = map(pe.PartialVal.unknown, in_avals)\n  jaxpr, out_pvals, const = pe.trace_to_jaxpr_nounits(flat_fun, in_pvals,\n                                                      instantiate=True)\n  jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), True)\n  out_avals, _ = unzip2(out_pvals)\n  out_dtypes = [a.dtype for a in out_avals if not a.is_high]\n  if not (all(dtypes.issubdtype(d, np.inexact) for d in in_dtypes + out_dtypes)\n          or all(dtypes.issubdtype(d, np.integer)\n                 for d in in_dtypes + out_dtypes)):","sourceCodeStart":2076,"sourceCodeEnd":2112,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L2076-L2112","documentation":"jax.linear_transpose used to accept a reduce_axes parameter for reductions across named axes; this was never fully supported and is now deprecated, so passing any truthy value raises NotImplementedError immediately.","triggerScenarios":"Calling jax.linear_transpose(f, x, y, reduce_axes=('batch',)) or any call that includes the reduce_axes keyword.","commonSituations":"Old code written against an early JAX version that allowed reduce_axes; copying examples from outdated docs or Stack Overflow answers.","solutions":["Remove the reduce_axes argument entirely","Handle reductions over batched axes yourself, e.g. by closing over fixed axis sizes or using vmap instead"],"exampleFix":"# before\nf_t = jax.linear_transpose(f, x, y, reduce_axes=('batch',))\n# after\nf_t = jax.linear_transpose(f, x, y)","handlingStrategy":"validation","validationCode":"jax.linear_transpose(f, x, y)  # simply never pass reduce_axes","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Search codebase for reduce_axes usages when upgrading JAX","Use vmap for batched transposes instead of reduce_axes"],"tags":["jax","linear-transpose","deprecated-api"],"backgroundTag":"deprecated-argument-removed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}