{"record":{"id":"81ead51d55116132","repo":"jax-ml/jax","slug":"open-an-issue-at-https-github-com-google-jax","errorCode":null,"errorMessage":"open an issue at https://github.com/google/jax !!","messagePattern":"open an issue at https://github\\.com/google/jax !!","errorType":"error_code","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/solves.py","lineNumber":409,"sourceCode":"\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\ndef _linear_solve_batching_rule(axis_data, args, dims, const_lengths, jaxprs):\n  orig_bat = [d is not None for d in dims]\n\n  params, b = _split_linear_solve_args(args, const_lengths)\n  params_dims, b_dims = _split_linear_solve_args(dims, const_lengths)","sourceCodeStart":391,"sourceCodeEnd":427,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/solves.py#L391-L427","documentation":"Inside the custom_linear_solve transpose rule, after splitting primals into params and b, the rule asserts that only the b (right-hand-side) arguments carry tangent/undefined-primal status. If any matrix/operator parameter itself is being differentiated through (an 'open' AD input in the transpose), the rule hits an unimplemented path and raises NotImplementedError asking for a GitHub issue.","triggerScenarios":"Taking jax.grad of a function where custom_linear_solve's matvec closure captures differentiable (traced) arrays — i.e. differentiating w.r.t. the linear operator's parameters, in a configuration not covered by the implicit-diff path (e.g. no transpose_solve dealing with params).","commonSituations":"End-to-end differentiation through a neural PDE solver or learned preconditioner used inside custom_linear_solve without providing transpose_solve; jvp works but vjp/grad hits this.","solutions":["Provide an explicit transpose_solve that handles the operator's parameters so the generic path is bypassed","Stop-gradient the operator parameters: jax.lax.stop_gradient on the matvec's captured weights if you only need grads w.r.t. b or upstream inputs","Reparameterize so the operator is constant: precompute the matrix outside the traced function","If it looks like a genuine unsupported case, open the requested issue at github.com/google/jax"],"exampleFix":"// before\nx = jax.lax.custom_linear_solve(lambda v: apply_A(A, v), b, solve)\njax.grad(f)(...)  # A differentiable -> transpose rule fails\n// after\nA_ = jax.lax.stop_gradient(A)\nx = jax.lax.custom_linear_solve(lambda v: apply_A(A_, v), b, solve)","handlingStrategy":"fallback","validationCode":"# stop-gradient operator params before the solve\nmatvec_safe = lambda v: matvec(jax.lax.stop_gradient(A), v)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Wrap operator parameters in stop_gradient when you don't need their grads","Provide transpose_solve handling params for full operator differentiation"],"tags":["jax","autodiff","custom-linear-solve","not-implemented"],"backgroundTag":"autodiff-unsupported-input","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}