# Copyright 2022 The JAX Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Module for the custom linear solve and utilities.""" from functools import partial import operator from typing import Any, NamedTuple from collections.abc import Callable from jax._src import ad_util from jax._src import api from jax._src import api_util from jax._src import core from jax._src import hijax from jax._src.interpreters import ad from jax._src.interpreters import batching from jax._src.interpreters import mlir from jax._src.interpreters import partial_eval as pe from jax._src.traceback_util import api_boundary from jax._src import flattree as ft from jax._src.tree_util import tree_leaves, tree_map, tree_structure from jax._src.util import split_list, safe_map import numpy as np from jax._src.lax.control_flow.common import ( _check_tree, ) _map = safe_map @api_boundary def custom_root(f: Callable, initial_guess: Any, solve: Callable[[Callable, Any], Any], tangent_solve: Callable[[Callable, Any], Any], has_aux=False): """Differentiably solve for the roots of a function. This is a low-level routine, mostly intended for internal use in JAX. Gradients of custom_root() are defined with respect to closed-over variables from the provided function ``f`` via the implicit function theorem: https://en.wikipedia.org/wiki/Implicit_function_theorem Args: f: function for which to find a root. Should accept a single argument, return a tree of arrays with the same structure as its input. initial_guess: initial guess for a zero of f. solve: function to solve for the roots of f. Should take two positional arguments, f and initial_guess, and return a solution with the same structure as initial_guess such that func(solution) = 0. In other words, the following is assumed to be true (but not checked):: solution = solve(f, initial_guess) error = f(solution) assert all(error == 0) tangent_solve: function to solve the tangent system. Should take two positional arguments, a linear function ``g`` (the function ``f`` linearized at its root) and a tree of array(s) ``y`` with the same structure as initial_guess, and return a solution ``x`` such that ``g(x)=y``: - For scalar ``y``, use ``lambda g, y: y / g(1.0)``. - For vector ``y``, you could use a linear solve with the Jacobian, if dimensionality of ``y`` is not too large: ``lambda g, y: np.linalg.solve(jacobian(g)(y), y)``. has_aux: bool indicating whether the ``solve`` function returns auxiliary data like solver diagnostics as a second argument. Returns: The result of calling solve(f, initial_guess) with gradients defined via implicit differentiation assuming ``f(solve(f, initial_guess)) == 0``. """ guess_avals = tree_map(core.typeof, initial_guess) guess_tree = tree_structure(initial_guess) f_traced = api.jit(f).trace(initial_guess) _check_tree("f", "initial_guess", f_traced.out_tree, guess_tree, False) f_consts, f = f_traced.closure_convert() solve_traced = api.jit(partial(solve, partial(f, f_consts))).trace(initial_guess) _check_tree("solve", "initial_guess", solve_traced.out_tree, guess_tree, has_aux) solve_consts, solve = solve_traced.closure_convert() # Capture tangent_solve's environment without constructing the derivative of f. tangent_avals = tree_map(lambda a: a.to_tangent_aval(), guess_avals) f_jvp = _RootLinearMap( f, (tree_map(core.typeof, f_consts), guess_avals, tangent_avals), tree_map(lambda a: a.to_tangent_aval(), f_traced.out_avals), f_traced.effects) def linearize_and_solve(x, b): return tangent_solve(partial(f_jvp, f_consts, x), b) tangent_traced = api.jit(linearize_and_solve).trace(guess_avals, tangent_avals) _check_tree("tangent_solve", "x", tangent_traced.out_tree, guess_tree, False) tangent_consts, tangent = tangent_traced.closure_convert() # Keep solve's operands first so its input effect indices carry over directly. args = (solve_consts, initial_guess, f_consts, tangent_consts) return CustomRoot(f, solve, tangent, solve_traced, tree_map(core.typeof, args), has_aux)(*args) class CustomRoot(hijax.HiPrim): f: Callable solve: Callable tangent_solve: Callable solve_jaxpr: core.Jaxpr has_aux: bool skip_linearization_on_zero_tangents = True def __init__(self, f, solve, tangent_solve, solve_traced, in_avals, has_aux): self.in_avals = in_avals self.out_aval = solve_traced.out_avals self.effects = solve_traced.effects self.params = dict(f=f, solve=solve, tangent_solve=tangent_solve, solve_jaxpr=solve_traced.jaxpr.replace(consts=None), has_aux=has_aux) super().__init__() def pp_params(self): # f, solve, and tangent_solve are closure-converted return dict(solve_jaxpr=self.solve_jaxpr, has_aux=self.has_aux) def check(self, *_): hijax.check_custom_effects(self.effects, 'custom_root') def expand(self, solve_consts, guess, f_consts, tangent_consts): # pyrefly: ignore[bad-override] return self.solve(solve_consts, guess) def lin(self, nzs, solve_consts, guess, f_consts, tangent_consts): # pyrefly: ignore[bad-override] sol = self(solve_consts, guess, f_consts, tangent_consts) if not any(tree_leaves(nzs[2])): return sol, None, False root, aux = sol if self.has_aux else (sol, ()) _, f_lin = api.linearize(lambda p: self.f(p, root), f_consts, in_nzs=(nzs[2],)) aux_dot = tree_map(ad_util.p2tz, aux) nzs_out = (True, False) if self.has_aux else True return sol, (root, tangent_consts, f_lin, aux_dot), nzs_out def linearized(self, res, solve_dot, guess_dot, f_dot, tangent_dot): # pyrefly: ignore[bad-override] if res is None: return tree_map(ad_util.a2tz, self.out_aval) root, tangent_consts, f_lin, aux_dot = res rhs = f_lin(tree_map(ad_util.instantiate, f_dot, is_leaf=lambda x: isinstance(x, ad_util.Zero))) # F(p, x(p)) = 0 implies dx = -(D_x F)^-1 D_p F dp. root_dot = tree_map(operator.neg, self.tangent_solve(tangent_consts, root, rhs)) return (root_dot, aux_dot) if self.has_aux else root_dot jvp = hijax.jvp_from_lin vjp_fwd, vjp_bwd_retval = hijax.vjp_from_lin def batch(self, axis_data, args, dims): solve_dims, guess_dims, f_dims, tangent_dims = dims is_none = lambda x: x is None in_dims = tree_leaves(dims[:2], is_leaf=is_none) _, out_dims = batching.batch_jaxpr2(self.solve_jaxpr, axis_data, tuple(in_dims)) out_dims = self.out_tree.unflatten(out_dims) root_dims = out_dims[0] if self.has_aux else out_dims if any(d is not None for d in tree_leaves(dims, is_leaf=is_none)): root_avals = self.out_aval[0] if self.has_aux else self.out_aval root_dims = tree_map(lambda a: a.leading_axis_spec(), root_avals) sol_dims = (root_dims, out_dims[1]) if self.has_aux else root_dims solve = hijax.vmap_rule(axis_data, self.solve, (solve_dims, guess_dims), sol_dims) f = hijax.vmap_rule(axis_data, self.f, (f_dims, root_dims), root_dims) tangent_solve = hijax.vmap_rule(axis_data, self.tangent_solve, (tangent_dims, root_dims, root_dims), root_dims) in_avals = hijax.unmap_avals(axis_data, self.in_avals, dims) solve_traced = api.jit(solve).trace(*in_avals[:2]) prim = CustomRoot(f, solve, tangent_solve, solve_traced, in_avals, self.has_aux) return prim(*args), sol_dims class _RootLinearMap(hijax.HiPrim): """Apply D_x f or its transpose, expanding AD only when the map is used. Keeping the map opaque also lets tangent_solve build a Jacobian with jacfwd or jacrev without differentiating f during the primal custom_root trace. """ f: Callable transposed: bool skip_linearization_on_zero_tangents = True def __init__(self, f, in_avals, out_aval, effects, transposed=False): self.in_avals = in_avals self.out_aval = out_aval self.params = dict(f=f, transposed=transposed) self.effects = effects super().__init__() def expand(self, consts, x, v): # pyrefly: ignore[bad-override] if self.transposed: return api.vjp(partial(self.f, consts), x)[1](v)[0] return api.jvp(partial(self.f, consts), (x,), (v,))[1] def transpose(self, out_ct, consts, x, v_accums): # pyrefly: ignore[bad-override] in_avals = (*self.in_avals[:2], tree_map(lambda a: a.to_ct_aval(), self.out_aval)) out_aval = tree_map(lambda a: a.to_ct_aval(), self.in_avals[2]) prim = _RootLinearMap(self.f, in_avals, out_aval, self.effects, not self.transposed) out_ct = tree_map(ad_util.instantiate, out_ct, is_leaf=lambda x: isinstance(x, ad_util.Zero)) v_ct = prim(consts, x, out_ct) for acc, ct in zip(tree_leaves(v_accums), tree_leaves(v_ct)): if isinstance(acc, ad.GradAccum): acc.accum(ct) def jvp(self, primals, tangents): const_dot, x_dot, v_dot = tangents zero = lambda x: isinstance(x, ad_util.Zero) if all(zero(t) for t in tree_leaves((const_dot, x_dot), is_leaf=zero)): consts, x, _ = primals out = self(*primals) if all(zero(t) for t in tree_leaves(v_dot, is_leaf=zero)): return out, tree_map(ad_util.p2tz, out) v_dot = tree_map(ad_util.instantiate, v_dot, is_leaf=zero) return out, self(consts, x, v_dot) # Preserve symbolic zeros when differentiating the operator's coefficients. primals_ft = ft.flatten(primals) tangents_ft = primals_ft.update(self.in_tree.flatten_up_to(tangents)) out, out_dot = ad.jvp(self.expand, primals_ft, tangents_ft, instantiate=False) return out.unflatten(), out_dot.unflatten() lin, linearized = hijax.linearize_from_jvp vjp_fwd, vjp_bwd_retval = hijax.vjp_from_lin def batch(self, axis_data, args, dims): if all(d is None for d in tree_leaves(dims, is_leaf=lambda x: x is None)): return self(*args), tree_map(lambda _: None, self.out_aval) consts, x, v = args const_dims, x_dims, v_dims = dims lead = lambda avals: tree_map(lambda a: a.leading_axis_spec(), avals) x_lead, v_lead, out_dims = lead(self.in_avals[1]), lead(self.in_avals[2]), lead(self.out_aval) align = lambda y, src, dst: hijax.vmap_rule(axis_data, lambda y: y, (src,), dst)(y) f = hijax.vmap_rule(axis_data, self.f, (const_dims, x_lead), v_lead if self.transposed else out_dims) in_avals = hijax.unmap_avals(axis_data, self.in_avals, (const_dims, x_lead, v_lead)) out_aval = hijax.unmap_avals(axis_data, self.out_aval, out_dims) prim = _RootLinearMap(f, in_avals, out_aval, self.effects, self.transposed) return prim(consts, align(x, x_dims, x_lead), align(v, v_dims, v_lead)), out_dims class _LinearSolveTuple(NamedTuple): matvec: Any vecmat: Any solve: Any transpose_solve: Any def transpose(self): return type(self)(self.vecmat, self.matvec, self.transpose_solve, self.solve) def _split_linear_solve_args(args, const_lengths): params_list = split_list(args, list(const_lengths)) return _LinearSolveTuple(*params_list[:-1]), params_list[-1] def _transpose_one_output(linear_fun, primals): transpose_fun = api.linear_transpose(linear_fun, primals) def transposed_fun(x): (y,) = transpose_fun(x) return y return transposed_fun def _flatten(args): return [x for arg in args for x in arg] def _check_shapes(func_name, expected_name, actual, expected): actual_shapes = _map(np.shape, actual) expected_shapes = _map(np.shape, expected) if actual_shapes != expected_shapes: raise ValueError( f"{func_name}() output shapes must match {expected_name}, " f"got {actual_shapes} and {expected_shapes}") @partial(api_boundary, repro_api_name="jax.custom_linear_solve") def custom_linear_solve( matvec: Callable, b: Any, solve: Callable[[Callable, Any], Any], transpose_solve: Callable[[Callable, Any], Any] | None = None, symmetric=False, has_aux=False): """Perform a matrix-free linear solve with implicitly defined gradients. This function allows for overriding or defining gradients for a linear solve directly via implicit differentiation at the solution, rather than by differentiating *through* the solve operation. This can sometimes be much faster or more numerically stable, or differentiating through the solve operation may not even be implemented (e.g., if ``solve`` uses ``lax.while_loop``). Required invariant:: x = solve(matvec, b) # solve the linear equation assert matvec(x) == b # not checked Args: matvec: linear function to invert. Must be differentiable. b: constant right handle side of the equation. May be any nested structure of arrays. solve: higher level function that solves for solution to the linear equation, i.e., ``solve(matvec, x) == x`` for all ``x`` of the same form as ``b``. This function need not be differentiable. transpose_solve: higher level function for solving the transpose linear equation, i.e., ``transpose_solve(vecmat, x) == x``, where ``vecmat`` is the transpose of the linear map ``matvec`` (computed automatically with autodiff). Required for backwards mode automatic differentiation, unless ``symmetric=True``, in which case ``solve`` provides the default value. symmetric: bool indicating if it is safe to assume the linear map corresponds to a symmetric matrix, i.e., ``matvec == vecmat``. has_aux: bool indicating whether the ``solve`` and ``transpose_solve`` functions return auxiliary data like solver diagnostics as a second argument. Returns: Result of ``solve(matvec, b)``, with gradients defined assuming that the solution ``x`` satisfies the linear equation ``matvec(x) == b``. """ if transpose_solve is None and symmetric: transpose_solve = solve b_flat = ft.flatten(b) b_avals = b_flat.map(core.typeof) tree = b_flat.tree def _shape_checked(fun, name, has_aux): def f(x): y = fun(x) _check_shapes(name, "b", tree_leaves(y), b_flat) return y def f_aux(x): y, aux = fun(x) _check_shapes(name, "b", tree_leaves(y), b_flat) return y, aux return f_aux if has_aux else f matvec_debug = api_util.debug_info("custom_linear_solve", matvec, (b,), {}) # no auxiliary data assumed for matvec args_avals = ft.pack(((b_avals,),{})) matvec_jaxpr, out_avals = pe.trace_to_jaxpr( _shape_checked(matvec, "matvec", False), args_avals, matvec_debug) matvec_jaxpr, matvec_consts = pe.separate_consts(matvec_jaxpr) _check_tree("matvec", "b", out_avals.tree, tree, False) solve_debug = api_util.debug_info("custom_linear_solve solve", solve, (matvec, b), {}, static_argnums=(0,)) solve_jaxpr, out_avals = pe.trace_to_jaxpr( _shape_checked(partial(solve, matvec), "solve", has_aux), args_avals, solve_debug) solve_jaxpr, solve_consts = pe.separate_consts(solve_jaxpr) _check_tree("solve", "b", out_avals.tree, tree, has_aux) if transpose_solve is None: vecmat_jaxpr = tr_solve_jaxpr = None vecmat_consts = tr_solve_consts = [] else: transpose_solve_debug = api_util.debug_info( "custom_linear_solve transpose_solve", transpose_solve, (matvec, b), {}, static_argnums=(0,)) if symmetric: vecmat = matvec vecmat_jaxpr = matvec_jaxpr vecmat_consts = matvec_consts else: vecmat = _transpose_one_output(matvec, b) vecmat_jaxpr, out_avals = pe.trace_to_jaxpr( vecmat, args_avals, transpose_solve_debug) vecmat_jaxpr, vecmat_consts = pe.separate_consts(vecmat_jaxpr) assert out_avals.tree == tree tr_solve_jaxpr, out_avals = pe.trace_to_jaxpr( _shape_checked(partial(transpose_solve, vecmat), "transpose_solve", has_aux), args_avals, transpose_solve_debug) tr_solve_jaxpr, tr_solve_consts = pe.separate_consts(tr_solve_jaxpr) _check_tree("transpose_solve", "b", out_avals.tree, tree, has_aux) all_consts = [matvec_consts, vecmat_consts, solve_consts, tr_solve_consts] const_lengths = _LinearSolveTuple(*_map(len, all_consts)) jaxprs = _LinearSolveTuple( matvec_jaxpr, vecmat_jaxpr, solve_jaxpr, tr_solve_jaxpr) args = _flatten(all_consts) + list(b_flat) args = core.auto_insert_reshard(*args) out_flat = linear_solve_p.bind(*args, const_lengths=const_lengths, jaxprs=jaxprs) return out_avals.update(out_flat).unflatten() def _linear_solve_abstract_eval(*args, const_lengths, jaxprs): args_to_raise = args[sum(const_lengths):] # raise aux_args to shaped arrays as well if present # number of aux args is the difference in out_avals # of solve and matvec (since they map to the same vector space) num_aux = len(jaxprs.solve.out_avals) - len(jaxprs.matvec.out_avals) if num_aux > 0: args_to_raise += tuple(jaxprs.solve.out_avals[-num_aux:]) out_vma = core.standard_vma_rule('linear_solve', *args_to_raise) if any(core.getu(a) or core.getr(a) for a in args_to_raise): raise NotImplementedError return (tuple(a.update(manual_axis_type=a.mat.update(varying=out_vma)) for a in args_to_raise), jaxprs.solve.effects) def _custom_linear_solve_impl(*args, const_lengths, jaxprs): params, b = _split_linear_solve_args(args, const_lengths) x = core.jaxpr_as_fun(jaxprs.solve)(*(params.solve + b)) return x def _tangent_linear_map(func: Callable, params, params_dot, debug_info: core.DebugInfo, *x): """Compute the tangent of a linear map. Assuming ``func(*params, *x)`` is linear in ``x`` and computes ``A @ x``, this function computes ``∂A @ x``. """ assert any(type(p) is not ad_util.Zero for p in params_dot) zeros = _map(ad_util.p2tz, x) primals_ft = ft.flatten_list(params + list(x)) tangents_ft = ft.flatten_list(params_dot + zeros) _, out_tangent = ad.jvp(func, primals_ft, tangents_ft) return list(out_tangent) def _custom_linear_solve_jvp(primals, tangents, const_lengths, jaxprs): # A x - b = 0 # ∂A x + A ∂x - ∂b = 0 # ∂x = A^{-1} (∂b - ∂A x) kwargs = dict(const_lengths=const_lengths, jaxprs=jaxprs) x = linear_solve_p.bind(*primals, **kwargs) params, _ = _split_linear_solve_args(primals, const_lengths) params_dot, b_dot = _split_linear_solve_args(tangents, const_lengths) num_x_leaves = len(b_dot) # x is a flat tree with possible aux values appended # since x_tree == b_tree == b_dot_tree, we can cut off # aux values with len info provided by b_dot tree here x_leaves, _ = split_list(x, [num_x_leaves]) if all(type(p) is ad_util.Zero for p in params_dot.matvec): # no need to evaluate matvec_tangents rhs = b_dot else: matvec_tangents = _tangent_linear_map( core.jaxpr_as_fun(jaxprs.matvec), params.matvec, params_dot.matvec, jaxprs.matvec.debug_info, *x_leaves) rhs = _map(ad.add_tangents, b_dot, _map(operator.neg, matvec_tangents)) x_dot = linear_solve_p.bind(*(_flatten(params) + rhs), **kwargs) # split into x tangents and aux tangents (these become zero) dx_leaves, daux_leaves = split_list(x_dot, [num_x_leaves]) daux_leaves = _map(ad_util.p2tz, daux_leaves) x_dot = dx_leaves + daux_leaves return x, x_dot def _linear_solve_transpose_rule(cotangent, *primals, const_lengths, jaxprs): if jaxprs.transpose_solve is None: raise TypeError('transpose_solve required for backwards mode automatic ' 'differentiation of custom_linear_solve') params, b = _split_linear_solve_args(primals, const_lengths) if any(ad.is_undefined_primal(x) for xs in params for x in xs): raise NotImplementedError("open an issue at https://github.com/google/jax !!") assert all(ad.is_undefined_primal(x) for x in b) # TODO(mattjj): why? x_cotangent, other_cotangents = split_list(cotangent, [len(b)]) if any(type(ct) is not ad_util.Zero for ct in other_cotangents): raise NotImplementedError("open an issue at https://github.com/google/jax !!") del other_cotangents x_cotangent_ = _map(ad_util.instantiate, x_cotangent) cotangent_b_full = linear_solve_p.bind( *_flatten(params.transpose()), *x_cotangent_, const_lengths=const_lengths.transpose(), jaxprs=jaxprs.transpose()) cotangent_b, _ = split_list(cotangent_b_full, [len(b)]) return [None] * sum(const_lengths) + cotangent_b def _linear_solve_batching_rule(axis_data, args, dims, const_lengths, jaxprs): orig_bat = [d is not None for d in dims] params, b = _split_linear_solve_args(args, const_lengths) params_dims, b_dims = _split_linear_solve_args(dims, const_lengths) params_bat, orig_b_bat = _split_linear_solve_args(orig_bat, const_lengths) (matvec, vecmat, solve, solve_t) = jaxprs (matvec_bat, vecmat_bat, solve_bat, solve_t_bat) = params_bat # number of operator out avals is assumed to be the same for matvec/vecmat num_operator_out_avals = len(matvec.out_avals) num_aux = len(solve.out_avals) - num_operator_out_avals # Fixpoint computation of which parts of x and b are batched; we need to # ensure this is consistent between all four jaxprs b_bat = orig_b_bat x_bat = [False] * len(solve.out_avals) for i in range(1 + len(orig_b_bat) + len(solve.out_avals)): # Apply vecmat and solve -> new batched parts of x solve_jaxpr_batched, solve_x_bat = batching.batch_jaxpr( solve, axis_data, solve_bat + b_bat, instantiate=x_bat) if vecmat is None: vecmat_jaxpr_batched = None x_bat_out = solve_x_bat else: vecmat_jaxpr_batched, vecmat_x_bat = batching.batch_jaxpr( vecmat, axis_data, vecmat_bat + b_bat, instantiate=b_bat) # batch all aux data by default x_bat_out = _map(operator.or_, vecmat_x_bat + [True] * num_aux, solve_x_bat) # keep a slice of only the linear operator part of solve's avals x_bat_noaux = x_bat_out[:num_operator_out_avals] # Apply matvec and solve_t -> new batched parts of b matvec_jaxpr_batched, matvec_b_bat = batching.batch_jaxpr( matvec, axis_data, matvec_bat + x_bat_noaux, instantiate=b_bat) if solve_t is None: solve_t_jaxpr_batched = None b_bat_out = _map(operator.or_, matvec_b_bat, orig_b_bat) else: solve_t_jaxpr_batched, solve_t_b_aux_bat = batching.batch_jaxpr( solve_t, axis_data, solve_t_bat + x_bat_noaux, instantiate=x_bat_out) assert len(solve_t_b_aux_bat) == len(orig_b_bat) + num_aux solve_t_b_bat, _ = split_list(solve_t_b_aux_bat, [len(orig_b_bat)]) b_bat_out = _map(lambda m, s, o: m or s or o, matvec_b_bat, solve_t_b_bat, orig_b_bat) if x_bat_out == x_bat and b_bat_out == b_bat: break else: x_bat = x_bat_out b_bat = b_bat_out else: assert False, "Fixedpoint not reached" batched_jaxprs = _LinearSolveTuple(matvec_jaxpr_batched, vecmat_jaxpr_batched, solve_jaxpr_batched, solve_t_jaxpr_batched) # Move batched axes to the front new_params = [ batching.moveaxis(x, d, 0) if d is not None and d != 0 else x for x, d in zip(_flatten(params), _flatten(params_dims)) ] # Broadcast out b if necessary new_b = [ batching.broadcast(x, axis_data.size, 0, axis_data.explicit_mesh_axis) if now_bat and not was_bat else batching.moveaxis(x, d, 0) if now_bat and d != 0 else x for x, d, was_bat, now_bat in zip(b, b_dims, orig_b_bat, b_bat) ] outs = linear_solve_p.bind( *(new_params + new_b), const_lengths=const_lengths, jaxprs=batched_jaxprs) out_dims = [0 if batched else None for batched in solve_x_bat] return outs, out_dims linear_solve_p = core.Primitive('custom_linear_solve') linear_solve_p.multiple_results = True linear_solve_p.def_impl(_custom_linear_solve_impl) linear_solve_p.def_effectful_abstract_eval(_linear_solve_abstract_eval) ad.primitive_jvps[linear_solve_p] = _custom_linear_solve_jvp mlir.register_lowering( linear_solve_p, mlir.lower_fun(_custom_linear_solve_impl, multiple_results=True)) ad.primitive_transposes[linear_solve_p] = _linear_solve_transpose_rule batching.fancy_primitive_batchers[linear_solve_p] = _linear_solve_batching_rule