{"record":{"id":"67881fd536fa3a1c","repo":"jax-ml/jax","slug":"transpose-solve-required-for-backwards-mode-automa","errorCode":null,"errorMessage":"transpose_solve required for backwards mode automatic differentiation of custom_linear_solve","messagePattern":"transpose_solve required for backwards mode automatic differentiation of custom_linear_solve","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/solves.py","lineNumber":404,"sourceCode":"        core.jaxpr_as_fun(jaxprs.matvec), params.matvec, params_dot.matvec,\n        jaxprs.matvec.debug_info, *x_leaves)\n    rhs = _map(ad.add_tangents, b_dot, _map(operator.neg, matvec_tangents))\n\n  x_dot = linear_solve_p.bind(*(_flatten(params) + rhs), **kwargs)\n\n  # split into x tangents and aux tangents (these become zero)\n  dx_leaves, daux_leaves = split_list(x_dot, [num_x_leaves])\n\n  daux_leaves = _map(ad_util.p2tz, daux_leaves)\n\n  x_dot = dx_leaves + daux_leaves\n\n  return x, x_dot\n\n\ndef _linear_solve_transpose_rule(cotangent, *primals, const_lengths, jaxprs):\n  if jaxprs.transpose_solve is None:\n    raise TypeError('transpose_solve required for backwards mode automatic '\n                    'differentiation of custom_linear_solve')\n\n  params, b = _split_linear_solve_args(primals, const_lengths)\n  if any(ad.is_undefined_primal(x) for xs in params for x in xs):\n    raise NotImplementedError(\"open an issue at https://github.com/google/jax !!\")\n  assert all(ad.is_undefined_primal(x) for x in b)  # TODO(mattjj): why?\n  x_cotangent, other_cotangents = split_list(cotangent, [len(b)])\n  if any(type(ct) is not ad_util.Zero for ct in other_cotangents):\n    raise NotImplementedError(\"open an issue at https://github.com/google/jax !!\")\n  del other_cotangents\n  x_cotangent_ = _map(ad_util.instantiate, x_cotangent)\n  cotangent_b_full = linear_solve_p.bind(\n      *_flatten(params.transpose()), *x_cotangent_,\n      const_lengths=const_lengths.transpose(), jaxprs=jaxprs.transpose())\n  cotangent_b, _ = split_list(cotangent_b_full, [len(b)])\n  return [None] * sum(const_lengths) + cotangent_b\n\n","sourceCodeStart":386,"sourceCodeEnd":422,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/solves.py#L386-L422","documentation":"jax.lax.custom_linear_solve differentiates through either (a) solving another system with the same matrix, or (b) explicitly calling a user-provided transpose_solve. The transpose rule requires jaxprs.transpose_solve; if custom_linear_solve was called without transpose_solve and matrix_symmetric=True was not applicable (so no transpose jaxpr was built), backward-mode differentiation raises TypeError.","triggerScenarios":"Calling jax.grad (or jax.vjp) on a function containing custom_linear_solve where neither transpose_solve was passed nor matrix_symmetric=True was set, and the solve function is not implicitly transposable.","commonSituations":"Wrapping external solvers (CuSPARSE, PETSc, scipy with differing transpose semantics) for use under grad; forward-mode (jvp) works but switching to loss.backward-style vjp exposes the missing transpose path.","solutions":["Pass transpose_solve to custom_linear_solve, e.g. transpose_solve=lambda Mt, v: solve_for(Mt.T if Mt is not None else None, v) if symmetric; otherwise a solver for the transposed system","Set matrix_symmetric=True if the operator is symmetric so the same solve serves as transpose_solve","Use jax.jacrev with a forward-over-reverse scheme or jvp if only first derivatives of the solve w.r.t. b are needed"],"exampleFix":"// before\nx = jax.lax.custom_linear_solve(matvec, b, solve=direct_solve)\nloss = jax.grad(f)(x)\n// after\nx = jax.lax.custom_linear_solve(matvec, b, solve=direct_solve,\n                                 transpose_solve=lambda Mt, v: direct_solve(Mt, v))\nloss = jax.grad(f)(x)","handlingStrategy":"fallback","validationCode":"# ensure transpose path exists before grad\nassert matrix_symmetric or transpose_solve is not None, \\\n    'grad requires transpose_solve or matrix_symmetric=True'","typeGuard":null,"tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept TypeError as e:\n    if 'transpose_solve required' in str(e):\n        f2 = remake_with_transpose_solve(f)\n        return jax.grad(f2)(x)\n    raise","preventionTips":["Always pass transpose_solve when the operator is not symmetric","Set matrix_symmetric=True for symmetric operators","Test jax.grad paths of custom solves in CI"],"tags":["jax","autodiff","custom-linear-solve","transpose"],"backgroundTag":"missing-transpose-solve-for-vjp","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}