Repository: google/jax

Issue:
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.


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).

Resolving patch:
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
--- 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)
+
+  def do_transpose(primals_in, cotangents_in):
+    # NOTE: This is passing in undefined primals in place of tangent arguments, but it
+    #       should all work out, because we're only computing the primal part here.
+    residuals = core.jaxpr_as_fun(primal_jaxpr)(*primals_in)[len(cotangents_in):]
+    # Now that we have a purely linear jaxpr, we can transpose it
+    cotangents_out = backward_pass(tangent_jaxpr.jaxpr, (), primals_in + residuals, cotangents_in)
+    # backward_pass 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)]
+
+  flat_args, in_tree_def = tree_flatten((primals_in, cotangents_in))
+  flat_do_transpose, out_tree = flatten_fun_nokwargs(lu.wrap_init(do_transpose), in_tree_def)
+  flat_cotangents_out = pe.remat_call_p.bind(flat_do_transpose, *flat_args, **params)
+  return tree_unflatten(out_tree(), flat_cotangents_out)
+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)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- 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))