Repository: jax-ml/jax
Pull request: Simplify handling of non-linear equations in backward_pass and fix remat
Previously, `backward_pass` has been generalized to be able to handle
non-linear computation in the body, but it could easily get confused
into doing unnecessary work only to throw it away later. Additionally, it
treated any call primitive embedded inside remat like remat itself,
which is obviously wrong.

This patch fixes both of those issues and simplifies a bunch of the code
at the same time. `backward_pass` now has an invariant that it only
deals with jaxprs containing linear equations alone, and becomes
a simple transposing interpreter again.

**Background on JVP vs linearization**

Ok, so why does this change actually fix the problem? It is important to
understand that JVP and linearization transforms are actually two
different things, even though we often identify them as one. Both take
in a function of type `a -> b`, but their ranges are different! JVP
returns a function of type `(a, T a) -> (b, T b)` while linearization
returns `a -> (b, T a --o T b)`. Note that the second type carries more
information, because we get a guarantee that (1) `b` does not depend on
`T a` and (2) the dependence of `T b` on `T a` is linear.

The reason why we usually treat them as equivalent, is that they can be
shown to be "isomorphic". If we take the output of linearization, we can
make it a JVP-like function using the following combinator:
```haskell
jvp f = \a ta -> let (b, lf) = linearize f in (b, lf ta)
```
More importantly for JAX, which doesn't have a linearization interpreter,
if we assume (1) and (2), linearization can be recovered in terms of jvp
as well:
```haskell
linearize f = \a -> let fjvp = jvp f in
                    partial_eval fjvp (Known a) Unknown
```
That is, if we have a mathematically correct JVP, then linearization is
simply partial evaluation with all primal values marked as known, and
all tangents treated as yet unknown values.

One important performance consideration is that for forward-mode AD we
really want to use the JVP formulation, which can interleave the computation
of primals and tangents, instead of sequencing them and increasing the memory
cost. On the other hand, transposition (necessary for VJPs!) can only be
applied to linear functions, and so it can't possibly work on the output
of JVP. It really can only be apply to the second output of the
linearization transform. Hence, we really care about both, but can we avoid
having two very similar implementations of (approximately) the same thing?
It seems that the answer is yes, because of the equivalence outlined above!

**If all this is so nice, then what's the problem?**

The problem is, of course, remat. Partial eval is able to thread the
known/unknown information correctly through regular call primitives, but
mind you, remat is no regular call primitive! Once we enter remat, we are
no longer interested in treating _anything_ like a known value. After
all, our goal here is to record an accurate trace of everything that has
happened in the body of a remat, including the primal (known!)
computation. This however presents a challenge for implementing
linearization in terms of JVP, because inside the body of remat we break
the assumption that known/unknown corresponds to the primal/tangent
distinction. Its body, instead of representing the second output of
linearization simply contains the traced JVP code now...

One way to fix it would be to implement a proper linearization pass that
would track the distinciton between primal and tangent information while
still allowing to stage out code for primals. @mattjj and I have even
started hacking together an implementation for that.

I've been trying to convince @mattjj that there is no other way to go
about it, but I couldn't really convince him that this is the case.
Then, once I wanted to write a semi-formal proof I could no longer even
convince myself! Turns out that there is an alternative solution!

What this patch does is, it stops caring about the output of the
`linearize` function (defined as JVP + partial eval, as discussed above)
to be a good linearization. It still is if you don't use remats in your
code, but it still breaks miserably once you do. However, as long as all
the complications are contained solely in the `call_jaxpr` embedded inside
a remat, we still have a chance to fix them! This is because the
transposition interpreter never reaches into those bodies directly, but
rather asks the call primitive to transpose itself.

Now, how do you transpose remat? We can't just reuse the code used for
regular call primitives (this is what happens now BTW), because unlike
for them, the `call_jaxpr` doesn't represent a linear function! But it's
not completely useless either --- it contains the traced JVP code. So,
how do we get from there to a linear function? Partial eval! And if you
think about it, it is exactly what we wanted --- we end up evaluating all
the primal code in the body once again, while only staging out the tangent
computation, to be passed into the transposing interpreter again.

Fin.

(Also, fixes #2737)

Problem statement:
remat + eager mode = unnecessary FLOPs in backward_pass
While `remat` takes great care to avoid unnecessary FLOPs during tracing of the primal code (see the discussion around the `concrete` parameter), it looks like some FLOPs have slipped through the cracks in the `backward_pass`! To verify you can apply the diff below, that I'm folding for readability.
<details>

```diff
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
index 0b730e0a..ba8bcf5c 100644
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -187,8 +188,14 @@ def backward_pass(jaxpr: core.Jaxpr, consts, args, cotangents_in):
       if any(is_linear(v) for v in eqn.invars):
         linear_eqns.append(eqn)
       if any(not is_linear(v) for v in eqn.invars):
+        print('Evaluating primal call primitive')
         ans = _eval_subjaxpr_primals(eqn.primitive, call_jaxpr,
                                      map(read_primal, eqn.invars), params)
+        print('Results:', ans)
         map(write_primal, eqn.outvars, ans)

   ct_env: Dict[Any, Any] = {}
@@ -220,6 +228,7 @@ def _eval_subjaxpr_primals(prim, jaxpr, in_vals, params):
   return tree_unflatten(out_tree(), out_flat)

 def _eval_primals(jaxpr, args):
+  print('eval_primals >>>>')
   primal_env = {}

   def read_primal(v):
@@ -230,6 +239,7 @@ def _eval_primals(jaxpr, args):

   def write_primal(v, val):
     if not is_undefined_primal(val):
+      print('COMPUTED: ', val)
       primal_env[v] = val

   def is_linear(var):
@@ -256,6 +266,7 @@ def _eval_primals(jaxpr, args):
         ans = _eval_subjaxpr_primals(eqn.primitive, call_jaxpr,
                                      map(read_primal, eqn.invars), params)
         map(write_primal, eqn.outvars, ans)
+  print('eval_primals <<<<')
   return map(read_primal, jaxpr.outvars)

 class UndefinedPrimal:
```

</details>

Then, running this code: 
```py
@checkpoint
def g(x):
    return x ** 3

def f(x):
    return g(x) * 2

vjp(f, x)[1](x)
```
produces the following output:
```
Evaluating primal call primitive
eval_primals >>>>
COMPUTED:  *
COMPUTED:  [2.]
COMPUTED:  [4.]
COMPUTED:  [12.]
eval_primals <<<<
Results: [*, UndefinedPrimal(ShapedArray(float32[1]))]
```
indicating that the call to `_eval_subjaxpr_primals` for the purpose of getting _potential primal outputs_ out of `remat_call` evaluated a bunch of expressions on real data, only to return a unit and an undefined primal. In effect, every `remat_call` gets evaluated twice, with all results of the first evaluation getting thrown away.

A potential solution would be to make `_eval_subjaxp_primals` only compute the primal values that actually affect the primal outputs (none in this case). The problem with that in that JVP jaxprs make no distinction between primals and tangents in their return lists, except that we might be able to depend on the convention that the first half of those always corresponds to primals, while the rest are their tangents. Not sure if that's a sound thing to do though (would that generalize to all call primitives?), so I wanted to discuss this before trying to resolve the issue. If yes, then `partial_eval._dce_jaxpr` might be a sufficient fix.

Hints from issue discussion:
Wow, brilliant catch!
I've forgotten the details now, but I think having some structure/bookkeeping to distinguish between primals and tangents from JVPs (and similarly for VJPs) would have been helpful when we were prototyping grad of sharded_jit (specifically in the sharded_jit translation rule I think). Not sure how relevant that is here but just throwing it out there.
@skye that's a really useful point! I think actually we might be able to collect a handful of such examples. (One that was on my mind recently: having linearity information available would make it much easier for users to write custom derivatives for higher-order functions that correctly handle closed-over traces in their function-valued arguments.)

@apaszke and I chatted this morning and, while there's still more thinking to do, it seems like it might not be a big lift to include this information in jaxprs, especially now that we include avals for each variable as a kind of type information. After all, when we differentiate things, we have the linearity information available; we just throw it away rather than recording it in the jaxpr. (There are a couple ad-hoc places we record it, like the `linear` params of `scan` and `cond`. But those feel analogous to the way we used to stash shape data in params, before we had avals available.)

cc @dougalm
Another interesting failure case that I forgot to document properly in here is that `call_p` primitive nested inside remat calls will end up behaving in the same way as if they were nested remat calls. To reproduce this you can use this code:
```python
def g(x):                                                           
    return jax.core.call(lu.wrap_init(lambda x: (np.sin(x),)), x)[0]

# g = remat(np.sin)
                                                                    
@remat                                                              
def f(x):                                                           
    return g(x * 2) * 4                                             
                                                                    
print(jax.make_jaxpr(lambda x, y: jax.jvp(f, [x], [y]))(x, x))      
```
which produces the following jaxpr:
```
{ lambda  ; a b.
  let c d = remat_call[ call_jaxpr={ lambda  ; a b.
                                     let c = mul a 2.0
                                         d = mul b 2.0
                                         e f = call[ call_jaxpr={ lambda  ; a b.
                                                                  let c = sin a
                                                                      d = cos a
                                                                      e = mul b d
                                                                  in (c, e) }
                                                     name=pe(jvp(<lambda>)) ] c d
                                         g = mul e 4.0
                                         h = mul f 4.0
                                     in (g, h) }
                        concrete=False
                        name=pe(jvp(f)) ] a b
  in (c, d) }
```

Uncommenting the line with `remat` produces an identical jaxpr, only with `call` replaced by `remat_call`. Since `remat` is (correctly) not special cased anywhere in AD, they will be treated in exactly the same way (they share their transpose implementation).

Commit message: Track the number of remat reevaluations in a test

Patch to review:
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
index cf52c42a1d2f..2e41f77168c1 100644
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -170,47 +170,18 @@ def write_primal(v, val):
   #        forces primal_in to contain UndefinedPrimals for tangent values!
   map(write_primal, jaxpr.invars, primals_in)
 
-  def is_linear(var):
-    if type(var) is Literal:
-      return False
-    else:
-      return var not in primal_env
-
-  linear_eqns = []
-  for eqn in jaxpr.eqns:
-    prim = eqn.primitive
-    if not (prim.call_primitive or prim.map_primitive):
-      if any(is_linear(v) for v in eqn.invars):
-        linear_eqns.append(eqn)
-      else:
-        in_vals = map(read_primal, eqn.invars)
-        ans = prim.bind(*in_vals, **eqn.params)
-        if prim.multiple_results:
-          map(write_primal, eqn.outvars, ans)
-        else:
-          write_primal(eqn.outvars[0], ans)
-    else:
-      call_jaxpr, params = core.extract_call_jaxpr(prim, eqn.params)
-      if any(is_linear(v) for v in eqn.invars):
-        linear_eqns.append(eqn)
-      if any(not is_linear(v) for v in eqn.invars):
-        # FIXME: Some invars correspond to tangents
-        ans = _eval_subjaxpr_primals(prim, call_jaxpr,
-                                     map(read_primal, eqn.invars), params)
-        map(write_primal, eqn.outvars, ans)
-
   # Find the last use of each cotangent so that they can be removed
   # as soon as possible.
   drop_cts: List[Set[Any]] = []
   seen_vars: Set[Any] = set(jaxpr.invars)
-  for eqn in linear_eqns:
+  for eqn in jaxpr.eqns:
     read_set = set(eqn.outvars)  # NOTE: eqn is not transposed yet!
     drop_cts.append(read_set - seen_vars)
     seen_vars |= read_set
 
   ct_env: Dict[Any, Any] = {}
   map(write_cotangent, jaxpr.outvars, cotangents_in)
-  for eqn, to_drop in zip(linear_eqns[::-1], drop_cts[::-1]):
+  for eqn, to_drop in zip(jaxpr.eqns[::-1], drop_cts[::-1]):
     # FIXME: Some invars correspond to tangents
     invals = map(read_primal, eqn.invars)
     if eqn.primitive.multiple_results:
@@ -218,9 +189,10 @@ def is_linear(var):
     else:
       cts_in, = map(read_cotangent, eqn.outvars)
     if eqn.primitive.call_primitive or eqn.primitive.map_primitive:
+      cts_in_avals = [v.aval for v in eqn.outvars]
       call_jaxpr, params = core.extract_call_jaxpr(eqn.primitive, eqn.params)
       cts_out = get_primitive_transpose(eqn.primitive)(
-          params, call_jaxpr, invals, cts_in)
+          params, call_jaxpr, invals, cts_in, cts_in_avals)
     else:
       cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals, **eqn.params)
     cts_out = [zero] * len(eqn.invars) if cts_out is zero else cts_out
@@ -232,60 +204,6 @@ def is_linear(var):
   cotangents_out = map(read_cotangent, jaxpr.invars)
   return cotangents_out
 
-def _eval_subjaxpr_primals(prim, jaxpr, in_vals, params):
-  assert not jaxpr.constvars
-  all_args, in_tree_def = tree_flatten((in_vals,))
-  fun = lu.hashable_partial(lu.wrap_init(_eval_primals), jaxpr)
-  fun, out_tree = flatten_fun_nokwargs(fun, in_tree_def)
-  assert prim.map_primitive ^ prim.call_primitive
-  if prim.map_primitive:
-    new_mapped_invars = [m for m, x in zip(params['mapped_invars'], in_vals)
-                         if not is_undefined_primal(x)]
-    new_params = dict(params, mapped_invars=tuple(new_mapped_invars))
-    out_flat = prim.bind(fun, *all_args, **new_params)
-  else:
-    out_flat = prim.bind(fun, *all_args, **params)
-  return tree_unflatten(out_tree(), out_flat)
-
-def _eval_primals(jaxpr, args):
-  primal_env = {}
-
-  def read_primal(v):
-    if type(v) is Literal:
-      return v.val
-    else:
-      return primal_env.get(v, UndefinedPrimal(v.aval))
-
-  def write_primal(v, val):
-    if not is_undefined_primal(val):
-      primal_env[v] = val
-
-  def is_linear(var):
-    if type(var) is Literal:
-      return False
-    else:
-      return var not in primal_env
-
-  write_primal(core.unitvar, core.unit)
-  assert not jaxpr.constvars
-  map(write_primal, jaxpr.invars, args)
-  for eqn in jaxpr.eqns:
-    if not (eqn.primitive.call_primitive or eqn.primitive.map_primitive):
-      if not any(is_linear(v) for v in eqn.invars):
-        in_vals = map(read_primal, eqn.invars)
-        ans = eqn.primitive.bind(*in_vals, **eqn.params)
-        if eqn.primitive.multiple_results:
-          map(write_primal, eqn.outvars, ans)
-        else:
-          write_primal(eqn.outvars[0], ans)
-    else:
-      call_jaxpr, params = core.extract_call_jaxpr(eqn.primitive, eqn.params)
-      if any(not is_linear(v) for v in eqn.invars):
-        ans = _eval_subjaxpr_primals(eqn.primitive, call_jaxpr,
-                                     map(read_primal, eqn.invars), params)
-        map(write_primal, eqn.outvars, ans)
-  return map(read_primal, jaxpr.outvars)
-
 class UndefinedPrimal:
   __slots__ = ['aval']
   def __init__(self, aval):
@@ -548,7 +466,7 @@ def traceable(num_primals, in_tree_def, *primals_and_tangents):
   yield out_flat, tree_def
 
 
-def call_transpose(primitive, params, call_jaxpr, args, ct):
+def call_transpose(primitive, params, call_jaxpr, args, ct, _):
   all_args, in_tree_def = tree_flatten(((), args, ct))  # empty consts
   fun = lu.hashable_partial(lu.wrap_init(backward_pass), call_jaxpr)
   fun, out_tree = flatten_fun_nokwargs(fun, in_tree_def)
@@ -556,9 +474,40 @@ def call_transpose(primitive, params, call_jaxpr, args, ct):
   out_flat = primitive.bind(fun, *all_args, **params)
   return tree_unflatten(out_tree(), out_flat)
 primitive_transposes[core.call_p] = partial(call_transpose, call_p)
-primitive_transposes[pe.remat_call_p] = partial(call_transpose, pe.remat_call_p)
 
-def map_transpose(primitive, params, call_jaxpr, args, ct):
+
+def remat_transpose(params, call_jaxpr, primals_in, cotangents_in, cotangent_in_avals):
+  # backward_pass can only transpose linear computations, but the call_jaxpr embedded in
+  # remat contains primal (non-linear) equations too. Hence, we have to eliminate those
+  # (in this case via partial_eval) before we call into backward_pass again.
+  typed_call_jaxpr = core.TypedJaxpr(
+      call_jaxpr, [],
+      [raise_to_shaped(p.aval if is_undefined_primal(p) else get_aval(p)) for p in primals_in],
+      cotangent_in_avals)
+  primal_jaxpr, tangent_jaxpr, out_unknowns = \
+    pe.partial_eval_jaxpr(typed_call_jaxpr,
+                          unknowns=map(is_undefined_primal, primals_in),
+                          instantiate=True,
+                          trace_type=None)
+
+  # Make a primal_jaxpr into a function. The flattening business is necessary so
+  # that we can thread UndefinedPrimals through the bind call. Note that jaxpr_as_fun
+  # doesn't understand how to handle those either, but tangent arguments are usued,
+  # so it should just ignore them.
+  primal_fun = lu.wrap_init(core.jaxpr_as_fun(primal_jaxpr))
+  flat_primals, in_tree_def = tree_flatten(primals_in)
+  primal_fun, _ = flatten_fun_nokwargs(primal_fun, in_tree_def)
+  residuals = pe.remat_call_p.bind(primal_fun, *flat_primals, **params)[len(cotangents_in):]
+
+  cotangents_out = call_transpose(pe.remat_call_p, params, tangent_jaxpr.jaxpr,
+                                  primals_in + residuals, cotangents_in, cotangent_in_avals)
+  # call_transpose will return cotangents computed for all invars, but some of them
+  # are residuals appended by partial eval, so we need to skip those before we return.
+  return cotangents_out[:len(primals_in)]
+primitive_transposes[pe.remat_call_p] = remat_transpose
+
+
+def map_transpose(primitive, params, call_jaxpr, args, ct, _):
   all_args, in_tree_def = tree_flatten(((), args, ct))  # empty consts
   fun = lu.hashable_partial(lu.wrap_init(backward_pass), call_jaxpr)
   fun, out_tree = flatten_fun_nokwargs(fun, in_tree_def)
diff --git a/tests/api_test.py b/tests/api_test.py
index 6602d376e81f..da858795d614 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1567,6 +1567,59 @@ def f(a_bool, y):
 
     api.jit(api.remat(f, concrete=True), static_argnums=0)(True, 1)  # no crash
 
+  def test_remat_eval_counter(self):
+    # https://github.com/google/jax/issues/2737
+    add_one_p = Primitive('add_one')
+    add_one = add_one_p.bind
+
+    num_evals = 0
+
+    @contextmanager
+    def assertEvals(n):
+      start = num_evals
+      yield
+      assert num_evals - start == n
+
+    def add_one_impl(x):
+      nonlocal num_evals
+      num_evals += 1
+      return x + 1
+    add_one_p.def_impl(add_one_impl)
+
+    def add_one_jvp(pin, tin):
+      pout = add_one(pin[0])
+      return pout, pout * tin[0]
+    ad.primitive_jvps[add_one_p] = add_one_jvp
+
+    add_one_p.def_abstract_eval(lambda x: x)
+
+    v = np.zeros((1,))
+
+    f = jax.remat(add_one)
+    g = jax.remat(lambda x: add_one(f(x)))
+
+    # 2 calls needed to evaluate g
+    with assertEvals(2):
+      _, vjp = jax.vjp(g, v)
+    # 2 calls made while transposing g, 1 call made while transposing f
+    with assertEvals(3):
+      vjp(v)
+
+    @jax.util.curry
+    def call(f, *args):
+      return jax.core.call(jax.linear_util.wrap_init(lambda *args: [f(*args)]), *args)[0]
+
+    f = call(add_one)
+    g = jax.remat(lambda x: add_one(f(x)))
+
+    # 2 calls needed to evaluate g
+    with assertEvals(2):
+      _, vjp = jax.vjp(g, v)
+    # 2 calls made while transposing g, no reevaluation for transposition of f
+    with assertEvals(2):
+      vjp(v)
+
+
   def test_trivial_computations(self):
     x = jnp.array([1, 2, 3])
     y = api.jit(lambda x: x)(x)
@@ -1721,7 +1774,7 @@ def fun(x):
     jaxpr = api.make_jaxpr(fun)(0.)
     self.assertMultiLineStrippedEqual("""
 { lambda b ; a.
-  let 
+  let
   in (a, 1.0, b) }
     """, str(jaxpr))
 


Code review comments:
diff_hunk: @@ -548,17 +466,50 @@ def traceable(num_primals, in_tree_def, *primals_and_tangents):
   yield out_flat, tree_def
 
 
-def call_transpose(primitive, params, call_jaxpr, args, ct):
+def call_transpose(primitive, params, call_jaxpr, args, ct, _):
   all_args, in_tree_def = tree_flatten(((), args, ct))  # empty consts
   fun = lu.hashable_partial(lu.wrap_init(backward_pass), call_jaxpr)
   fun, out_tree = flatten_fun_nokwargs(fun, in_tree_def)
   params = dict(params, name=wrap_name(params['name'], 'transpose'))
   out_flat = primitive.bind(fun, *all_args, **params)
   return tree_unflatten(out_tree(), out_flat)
 primitive_transposes[core.call_p] = partial(call_transpose, call_p)
-primitive_transposes[pe.remat_call_p] = partial(call_transpose, pe.remat_call_p)
 
-def map_transpose(primitive, params, call_jaxpr, args, ct):
+
+def remat_transpose(params, call_jaxpr, primals_in, cotangents_in, cotangent_in_avals):
+  # backward_pass can only transpose linear computations, but the call_jaxpr embedded in
+  # remat contains primal (non-linear) equations too. Hence, we have to eliminate those
+  # (in this case via partial_eval) before we call into backward_pass again.
+  typed_call_jaxpr = core.TypedJaxpr(
+      call_jaxpr, [],
+      [raise_to_shaped(p.aval if is_undefined_primal(p) else get_aval(p)) for p in primals_in],
+      cotangent_in_avals)
+  primal_jaxpr, tangent_jaxpr, out_unknowns = \
+    pe.partial_eval_jaxpr(typed_call_jaxpr,
+                          unknowns=map(is_undefined_primal, primals_in),
+                          instantiate=True,
+                          trace_type=None)
+
+  # Make a primal_jaxpr into a function. The flattening business is necessary so
+  # that we can thread UndefinedPrimals through the bind call. Note that jaxpr_as_fun
+  # doesn't understand how to handle those either, but tangent arguments are usued,
+  # so it should just ignore them.
+  primal_fun = lu.wrap_init(core.jaxpr_as_fun(primal_jaxpr))
+  flat_primals, in_tree_def = tree_flatten(primals_in)
+  primal_fun, _ = flatten_fun_nokwargs(primal_fun, in_tree_def)
+
+  params = dict(params, name=wrap_name(params['name'], 'transpose'))
+  residuals = pe.remat_call_p.bind(primal_fun, *flat_primals, **params)[len(cotangents_in):]
path: jax/interpreters/ad.py
text: @user2:
I adapted the above example like this:

```python
import jax
import jax.numpy as jnp

@user1.remat
def sinsin(x):
  return jnp.sin(jnp.sin(x))

def linearized(a, Ta):
  b, linearized = jax.linearize(sinsin, a)
  Tb = linearized(Ta)
  return b, Tb

print(jax.make_jaxpr(linearized)(1., 2.))
```

Here's the full jaxpr we get out:

```
{ lambda  ; a b.
  let c = sin a
      d = sin c
      e f = remat_call[ call_jaxpr={ lambda  ; a b.
                                     let c = cos a
                                         d = mul b c
                                         e = sin a
                                         f = cos e
                                         g = mul d f
                                     in (*, g) }
                        concrete=False
                        name=pe(pe(jvp(sinsin))) ] a b
  in (d, f) }
```

Looks great! Here's just the linear part of type `c -> T a --o T b`:

```
{ lambda b ; a.
  let c d = remat_call[ call_jaxpr={ lambda  ; a b.
                                     let c = cos a
                                         d = mul b c
                                         e = sin a
                                         f = cos e
                                         g = mul d f
                                     in (*, g) }
                        concrete=False
                        name=pe(pe(jvp(sinsin))) ] b a
  in (d,) }
```

I _think_ the ideal transpose of that linear part, of type `c -> CT b --o CT a`, might be something like (up to re-toposorting):

```
{ lambda b ; a.
  let c d = remat_call[ call_jaxpr={ lambda  ; a b.
                                     let c = cos a
                                         d = sin a
                                         e = cos d
                                         f = mul b e
                                         g = mul f c
                                     in (*, g) }
                        concrete=False
                        name=...) ] b a
  in (d,) }
```

In words, we still get one remat call but we have the linear part of it reversed. I think that if we then transposed again we'd be back where we started (up to re-toposorting).

I think this would look like the remat transpose rule first doing some jaxpr munging (essentially what it does now, using call_transpose or something similar) but to build a single jaxpr, and then ultimately to bind a single remat_call with that jaxpr.

WDYT? Setting aside _how_ we might do it, do you think the transpose jaxpr above looks like what we'd want, or do we want it to be different somehow?