jax-ml/jax · error · TypeError

linear_transpose only supports [float or complex] -> [float

Error message

linear_transpose only supports [float or complex] -> [float or complex], and integer -> integer functions, but got {in_dtypes} -> {out_dtypes}.

What it means

jax.linear_transpose mathematically requires the function to map floats/complex to floats/complex, or integers to integers. If input and output dtypes mix categories (e.g. float -> int or int -> float), no well-defined transpose exists and JAX raises TypeError.

Source

Thrown at jax/_src/api.py:2113

  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)):
    raise TypeError("linear_transpose only supports [float or complex] -> "
                    "[float or complex], and integer -> integer functions, "
                    f"but got {in_dtypes} -> {out_dtypes}.")

  @api_boundary
  def transposed_fun(const, out_cotangent):
    out_cts, out_tree2 = tree_flatten(out_cotangent)
    if out_tree() != out_tree2:
      raise TypeError("cotangent tree does not match function output, "
                      f"expected {out_tree()} but got {out_tree2}")
    if not all(map(core.typecheck, out_avals, out_cts)):
      raise TypeError("cotangent type does not match function output, "
                      f"expected {out_avals} but got {out_cts}")
    dummies = [ad.UndefinedPrimal(a.to_ct_aval()) for a in in_avals]
    in_cts = ad.backward_pass(jaxpr, True, const, dummies, out_cts)
    in_cts = map(ad.instantiate_zeros, in_cts)
    return tree_unflatten(in_tree, in_cts)

  # Ensure that transposed_fun is a PyTree

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove integer casts from the transposed function, or cast back to float at the end
2. Keep the whole pipeline in float/complex (or entirely integer)
3. Use jax.vjp or jax.grad instead if you actually want gradients of a float->int-mixed function (note: non-differentiable ops give zero cotangents)

Example fix

# before
f_t = jax.linear_transpose(lambda x: x.astype(jnp.int32), x)
# after
f_t = jax.linear_transpose(lambda x: x * 2.0, x)

import jax.numpy as jnp, numpy as np
from jax import dtypes
def transposable(f, *xs):
    dtypes_in = [x.dtype for x in xs]
    outs = jax.make_jaxpr(f)(*xs)[1]
    dts = [o.dtype for o in outs]
    ok = (all(dtypes.issubdtype(d, np.inexact) for d in dtypes_in + dts)
          or all(dtypes.issubdtype(d, np.integer) for d in dtypes_in + dts))
    return ok

When it happens

Trigger: Transposing f = lambda x: jnp.astype(x, jnp.int32) with float32 inputs; any function whose jaxpr mixes float inputs with integer outputs or vice versa.

Common situations: Index-generating functions (argmax, range-like ops) transposed for gradient purposes; functions that internally cast with astype to int; using linear_transpose where grad/vjp is intended.

Related errors

• cotangent type does not match function output, expected {out
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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