{"record":{"id":"2b7b17d95cc9837a","repo":"jax-ml/jax","slug":"linear-transpose-only-supports-float-or-complex","errorCode":null,"errorMessage":"linear_transpose only supports [float or complex] -> [float or complex], and integer -> integer functions, but got {in_dtypes} -> {out_dtypes}.","messagePattern":"linear_transpose only supports \\[float or complex\\] -> \\[float or complex\\], and integer -> integer functions, but got (.+?) -> (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":2113,"sourceCode":"  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)):\n    raise TypeError(\"linear_transpose only supports [float or complex] -> \"\n                    \"[float or complex], and integer -> integer functions, \"\n                    f\"but got {in_dtypes} -> {out_dtypes}.\")\n\n  @api_boundary\n  def transposed_fun(const, out_cotangent):\n    out_cts, out_tree2 = tree_flatten(out_cotangent)\n    if out_tree() != out_tree2:\n      raise TypeError(\"cotangent tree does not match function output, \"\n                      f\"expected {out_tree()} but got {out_tree2}\")\n    if not all(map(core.typecheck, out_avals, out_cts)):\n      raise TypeError(\"cotangent type does not match function output, \"\n                      f\"expected {out_avals} but got {out_cts}\")\n    dummies = [ad.UndefinedPrimal(a.to_ct_aval()) for a in in_avals]\n    in_cts = ad.backward_pass(jaxpr, True, const, dummies, out_cts)\n    in_cts = map(ad.instantiate_zeros, in_cts)\n    return tree_unflatten(in_tree, in_cts)\n\n  # Ensure that transposed_fun is a PyTree","sourceCodeStart":2095,"sourceCodeEnd":2131,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L2095-L2131","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Remove integer casts from the transposed function, or cast back to float at the end","Keep the whole pipeline in float/complex (or entirely integer)","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)"],"exampleFix":"# before\nf_t = jax.linear_transpose(lambda x: x.astype(jnp.int32), x)\n# after\nf_t = jax.linear_transpose(lambda x: x * 2.0, x)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp, numpy as np\nfrom jax import dtypes\ndef transposable(f, *xs):\n    dtypes_in = [x.dtype for x in xs]\n    outs = jax.make_jaxpr(f)(*xs)[1]\n    dts = [o.dtype for o in outs]\n    ok = (all(dtypes.issubdtype(d, np.inexact) for d in dtypes_in + dts)\n          or all(dtypes.issubdtype(d, np.integer) for d in dtypes_in + dts))\n    return ok","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid integer casts inside functions you plan to transpose","Run make_jaxpr and inspect dtypes before calling linear_transpose"],"tags":["jax","linear-transpose","dtype"],"backgroundTag":"unsupported-dtype-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}