jax-ml/jax · error · NotImplementedError

reduce_axes argument to transpose is deprecated

Error message

reduce_axes argument to transpose is deprecated

What it means

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.

Source

Thrown at jax/_src/api.py:2094

      See below for an example. (Note that the duck-typed objects cannot be
      namedtuples because those are treated as standard Python containers.)

  Returns:
    A callable that calculates the transpose of ``fun``. Valid input into this
    function must have the same shape/dtypes/structure as the result of
    ``fun(*primals)``. Output will be a tuple, with the same
    shape/dtypes/structure as ``primals``.

  >>> import jax
  >>>
  >>> f = lambda x, y: 0.5 * x - 0.5 * y
  >>> scalar = jax.ShapeDtypeStruct(shape=(), dtype=np.dtype(np.float32))
  >>> f_transpose = jax.linear_transpose(f, scalar, scalar)
  >>> f_transpose(1.0)
  (Array(0.5, dtype=float32), Array(-0.5, dtype=float32))
  """
  if reduce_axes:
    raise NotImplementedError("reduce_axes argument to transpose is deprecated")
  del reduce_axes
  primals_flat, in_tree = tree_flatten(primals)
  flat_fun, out_tree = flatten_fun_nokwargs(
      lu.wrap_init(fun,
                   debug_info=debug_info("linear_transpose", fun, primals, {})),
      in_tree)
  in_avals = [shaped_abstractify(x) for x in primals_flat]
  in_dtypes = [a.dtype for a in in_avals if not a.is_high]

  in_pvals = map(pe.PartialVal.unknown, in_avals)
  jaxpr, out_pvals, const = pe.trace_to_jaxpr_nounits(flat_fun, in_pvals,
                                                      instantiate=True)
  jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), True)
  out_avals, _ = unzip2(out_pvals)
  out_dtypes = [a.dtype for a in out_avals if not a.is_high]
  if not (all(dtypes.issubdtype(d, np.inexact) for d in in_dtypes + out_dtypes)
          or all(dtypes.issubdtype(d, np.integer)
                 for d in in_dtypes + out_dtypes)):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the reduce_axes argument entirely
2. Handle reductions over batched axes yourself, e.g. by closing over fixed axis sizes or using vmap instead

Example fix

# before
f_t = jax.linear_transpose(f, x, y, reduce_axes=('batch',))
# after
f_t = jax.linear_transpose(f, x, y)

jax.linear_transpose(f, x, y)  # simply never pass reduce_axes

When it happens

Trigger: Calling jax.linear_transpose(f, x, y, reduce_axes=('batch',)) or any call that includes the reduce_axes keyword.

Common situations: Old code written against an early JAX version that allowed reduce_axes; copying examples from outdated docs or Stack Overflow answers.

Related errors

• linear_transpose only supports [float or complex] -> [float
• cotangent tree does not match function output, expected {out
• cotangent type does not match function output, expected {out
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64

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