Repository: google/jax

Issue:
Adding `remat` causes an error to otherwise working code
When running the code below under `remat` we get the following error:

```
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-16-3143f2d8c2b3> in <module>()
     11 x = jnp.ones([1, 1, 1])
     12 f = jax.remat(f)
---> 13 jax.grad(f)(w, x)

17 frames
google3/third_party/py/jax/interpreters/ad.py in bilinear_transpose(lhs_rule, rhs_rule, cotangent, x, y, **kwargs)
    503 
    504 def bilinear_transpose(lhs_rule, rhs_rule, cotangent, x, y, **kwargs):
--> 505   assert (x is undefined_primal) ^ (y is undefined_primal)
    506   if x is undefined_primal:
    507     out = zero if cotangent is zero else lhs_rule(cotangent, y, **kwargs)

AssertionError: 
```

Without `jax.remat` the code does not error out and produces:

```
DeviceArray([[3.]], dtype=float32)
```

**Reproducer:**

```python
import jax
import jax.numpy as jnp

def f(w, x):
  a = jnp.dot(x, w)
  b = jnp.einsum("btd,bTd->btT", a, a)
  c = jnp.einsum("btT,btd->btd", b, a)
  return jnp.sum(c)

w = jnp.ones([1, 1])
x = jnp.ones([1, 1, 1])
f = jax.remat(f)  # comment this and things work
jax.grad(f)(w, x)
```
remat with scan part 2
This is a followup to #1907.

The code from #1907 does work now with the fix, but I am still getting the NotImplemented error I referred. The minimal example (below) is more complicated, but I hope it's not too bad. Please let me know if I can be helpful with debugging. 

Thanks in advance!


``` python
!pip install --upgrade -q https://storage.googleapis.com/jax-releases/cuda$(echo $CUDA_VERSION | sed -e 's/\.//' -e 's/\..*//')/jaxlib-$(pip search jaxlib | grep -oP '[0-9\.]+' | head -n 1)-cp36-none-linux_x86_64.whl
!pip install --upgrade -q jax

import numpy as onp

import jax.numpy as np
from jax import random, jit, grad, vmap, lax, remat
from functools import partial


#######################################
#Error only when all of these are True#
#######################################
function_test = True
test_remat = True
test_scan = True


def pot(dr, D0=5.0, **kwargs):
  U = dr**2
  if(function_test):
    U = D0 * U
  else:
    U = U - D0
  return np.array(U, dtype=dr.dtype)

def sys_energy(metric, D0=5.0):
  def fn_mapped(R, **dynamic_kwargs):
    _metric = vmap(vmap(partial(metric, **dynamic_kwargs), (0, None), 0), (None, 0), 0)
    dr = _metric(R, R)
    return np.sum(
        pot(dr, D0=D0),
        axis=None,
        keepdims=False)
  return fn_mapped

def create_simulation(params):
  D0 = params
  def displacement(Ra, Rb, **unused_kwargs):
    return Ra - Rb
  def shift(R, dR, **unused_kwargs):
    return R + 0.001*dR

  energy_fn = sys_energy(displacement, D0=D0)

  def get_dynamics_functions(energy_fn, shift):
    force_fn = grad(lambda R, *args, **kwargs: -energy_fn(R, *args, **kwargs))
    def init_fn(R):
      return R
    def apply_fn(state, **kwargs):
      R = state
      F = force_fn(R, **kwargs)
      R = shift(R, F)
      return R
    return init_fn, apply_fn

  return get_dynamics_functions(energy_fn, shift)

def run(params):
  key2, split = random.split(key,2)
  init_fn, apply_fn = create_simulation(params)

  R = random.uniform(
    split, (N, dimension), minval=0.0, maxval=8.0, dtype=np.float32)

  state = init_fn(R)

  n_steps_inner = 3
  n_steps_outer = 2

  def run_partial(state, i):
    def loop_fn(s, i):
      return apply_fn(s), np.array([0.])
    if(test_scan):
      steps = np.arange(0,n_steps_inner)
      return lax.scan(loop_fn, state, steps)
    else:
      return loop_fn(loop_fn(loop_fn(state,0)[0],0)[0],0)
    

  if(test_remat):
    run_partial = remat(run_partial)

  run_partial = jit(run_partial)

  #Surprisingly, this scan doesn't seem to affect things one way or another...
  state, temp = lax.scan(run_partial, state, np.arange(n_steps_outer))
  #state, temp = run_partial(run_partial(state,0)[0],0)
  

  return np.mean(state)


key = random.PRNGKey(0)
N = 10
dimension = 3

grad(run)(20.0)
```



Discussion:
@mattjj Does this reproduce for you?
Yes, thanks for the easy repro!
Here's the traceback, which ends somewhere inside `_scan_transpose`:
```
google3/third_party/py/jax/api.py in grad_f(*args, **kwargs)
    352   @wraps(fun, docstr=docstr, argnums=argnums)
    353   def grad_f(*args, **kwargs):
--> 354     _, g = value_and_grad_f(*args, **kwargs)
    355     return g
    356 

google3/third_party/py/jax/api.py in value_and_grad_f(*args, **kwargs)
    417              "differentiation, pass holomorphic=True.")
    418       raise TypeError(msg.format(dtype))
--> 419     g = vjp_py(onp.ones((), dtype=dtype))
    420     g = g[0] if isinstance(argnums, int) else g
    421     if not has_aux:

google3/third_party/py/jax/api.py in _vjp_pullback_wrapper(fun, cotangent_dtypes, io_tree, py_args)
   1237              "match type of corresponding primal output ({})")
   1238       raise TypeError(msg.format(_dtype(a), dtype))
-> 1239   ans = fun(*args)
   1240   return tree_unflatten(out_tree, ans)
   1241 

google3/third_party/py/jax/interpreters/ad.py in vjp_(*cts)
    112     dummy_primals_and_cts = (core.unit,) * len(cts) + cts
    113     dummy_args = (undefined_primal,) * len(jaxpr.invars)
--> 114     _, arg_cts = backward_pass(jaxpr, consts, (), dummy_args, dummy_primals_and_cts)
    115     arg_cts = arg_cts[len(primals):]
    116     return map(instantiate_zeros, primals, arg_cts)

google3/third_party/py/jax/interpreters/ad.py in backward_pass(jaxpr, consts, freevar_vals, args, cotangents_in)
    220       map(write_cotangent, bound_vars, ct_free_vars_out)
    221     else:
--> 222       cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals, **eqn.params)
    223     cts_out = [zero] * len(eqn.invars) if cts_out is zero else cts_out
    224     map(write_cotangent, eqn.invars, cts_out)

google3/third_party/py/jax/lax/lax_control_flow.py in _scan_transpose(cts, *args, **kwargs)
    831   #       jaxpr :: [ires, T d] -> [T c] -> [T a, eres] -> ([T c], [T b])
    832   # jaxpr_trans :: [ires] -> [CT d, CT c] -> [CT b, eres] -> ([CT d, CT c], [CT a])
--> 833   jaxpr_trans = _transpose_jaxpr(num_ires, num_consts - num_ires, num_eres, jaxpr)
    834   linear_trans = ([False] * num_ires +
    835                   [True] * (len(ct_consts) + len(ct_carry) + len(ct_ys)) +

google3/third_party/py/jax/lax/lax_control_flow.py in _transpose_jaxpr(num_res1, num_c, num_res2, jaxpr)
    864                 _map(ad.add_tangents, c_bar, new_c_bar))
    865     return c_bar + a_bar
--> 866   return _make_typed_jaxpr(transposed, res1_avals + c_avals + b_avals + res2_avals)
    867 
    868 def _make_typed_jaxpr(traceable, in_avals):

google3/third_party/py/jax/lax/lax_control_flow.py in _make_typed_jaxpr(traceable, in_avals)
    868 def _make_typed_jaxpr(traceable, in_avals):
    869   pvals = [pe.PartialVal((aval, core.unit)) for aval in in_avals]
--> 870   jaxpr, pvals_out, consts = pe.trace_to_jaxpr(traceable, pvals, instantiate=True)
    871   out_avals, _ = unzip2(pvals_out)
    872   return core.TypedJaxpr(jaxpr, consts, in_avals, out_avals)

google3/third_party/py/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out_calls)
    330   with new_master(trace_type) as master:
    331     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 332     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    333     assert not env
    334     del master

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    150     gen = None
    151 
--> 152     ans = self.f(*args, **dict(self.params, **kwargs))
    153     del args
    154     while stack:

google3/third_party/py/jax/lax/lax_control_flow.py in transposed(*res1_cbar_bbar_res2)
    858     primals = res1 + [ad.undefined_primal] * (num_c + num_a) + res2
    859     _, cbar_abar = ad.backward_pass(jaxpr.jaxpr, jaxpr.literals, (), primals,
--> 860                                     b_bar)
    861     _, new_c_bar, a_bar, _ = split_list(cbar_abar, [num_res1, num_c, num_a])
    862     a_bar = _map(ad.instantiate_zeros_aval, a_avals, a_bar)

google3/third_party/py/jax/interpreters/ad.py in backward_pass(jaxpr, consts, freevar_vals, args, cotangents_in)
    217       sub_freevar_vals = map(read_primal, bound_vars)
    218       ct_free_vars_out, cts_out = get_primitive_transpose(eqn.primitive)(
--> 219           eqn.params, subjaxpr, sub_consts, sub_freevar_vals, invals, cts_in)
    220       map(write_cotangent, bound_vars, ct_free_vars_out)
    221     else:

google3/third_party/py/jax/interpreters/ad.py in call_transpose(primitive, params, jaxpr, consts, freevar_vals, args, ct)
    548   fun = lu.hashable_partial(lu.wrap_init(backward_pass), jaxpr)
    549   fun, out_tree = flatten_fun_nokwargs(fun, in_tree_def)
--> 550   out_flat = primitive.bind(fun, *all_args, **params)
    551   return tree_unflatten(out_tree(), out_flat)
    552 primitive_transposes[core.call_p] = partial(call_transpose, call_p)

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **params)
    606   else:
    607     tracers = map(top_trace.full_raise, args)
--> 608     outs = map(full_lower, top_trace.process_call(primitive, f, tracers, params))
    609   return apply_todos(env_trace_todo(), outs)
    610 

google3/third_party/py/jax/interpreters/partial_eval.py in process_call(self, call_primitive, f, tracers, params)
    124     in_pvs, in_consts = unzip2([t.pval for t in tracers])
    125     fun, aux = partial_eval(f, self, in_pvs)
--> 126     out_flat = call_primitive.bind(fun, *in_consts, **params)
    127     out_pvs, jaxpr, env = aux()
    128     out_pv_consts, consts = split_list(out_flat, [len(out_flat)-len(jaxpr.constvars)])

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **params)
    603   if top_trace is None:
    604     with new_sublevel():
--> 605       outs = primitive.impl(f, *args, **params)
    606   else:
    607     tracers = map(top_trace.full_raise, args)

google3/third_party/py/jax/interpreters/xla.py in _xla_call_impl(fun, *args, **params)
    441   device = params['device']
    442   backend = params['backend']
--> 443   compiled_fun = _xla_callable(fun, device, backend, *map(arg_spec, args))
    444   try:
    445     return compiled_fun(*args)

google3/third_party/py/jax/linear_util.py in memoized_fun(fun, *args)
    221       fun.populate_stores(stores)
    222     else:
--> 223       ans = call(fun, *args)
    224       cache[key] = (ans, fun.stores)
    225     return ans

google3/third_party/py/jax/interpreters/xla.py in _xla_callable(fun, device, backend, *arg_specs)
    458   pvals = [pe.PartialVal((aval, core.unit)) for aval in abstract_args]
    459   with core.new_master(pe.StagingJaxprTrace, True) as master:
--> 460     jaxpr, (pvals, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
    461     assert not env  # no subtraces here
    462     del master, env

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    150     gen = None
    151 
--> 152     ans = self.f(*args, **dict(self.params, **kwargs))
    153     del args
    154     while stack:

google3/third_party/py/jax/interpreters/ad.py in backward_pass(jaxpr, consts, freevar_vals, args, cotangents_in)
    217       sub_freevar_vals = map(read_primal, bound_vars)
    218       ct_free_vars_out, cts_out = get_primitive_transpose(eqn.primitive)(
--> 219           eqn.params, subjaxpr, sub_consts, sub_freevar_vals, invals, cts_in)
    220       map(write_cotangent, bound_vars, ct_free_vars_out)
    221     else:

google3/third_party/py/jax/interpreters/ad.py in call_transpose(primitive, params, jaxpr, consts, freevar_vals, args, ct)
    548   fun = lu.hashable_partial(lu.wrap_init(backward_pass), jaxpr)
    549   fun, out_tree = flatten_fun_nokwargs(fun, in_tree_def)
--> 550   out_flat = primitive.bind(fun, *all_args, **params)
    551   return tree_unflatten(out_tree(), out_flat)
    552 primitive_transposes[core.call_p] = partial(call_transpose, call_p)

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **params)
    606   else:
    607     tracers = map(top_trace.full_raise, args)
--> 608     outs = map(full_lower, top_trace.process_call(primitive, f, tracers, params))
    609   return apply_todos(env_trace_todo(), outs)
    610 

google3/third_party/py/jax/interpreters/partial_eval.py in process_call(self, call_primitive, f, tracers, params)
    119       tracers = map(self.instantiate_const_abstracted, tracers)
    120     if call_primitive in call_partial_eval_rules:
--> 121       return call_partial_eval_rules[call_primitive](self, f, tracers, params)
    122     if call_primitive in map_primitives:
    123       return self.process_map(call_primitive, f, tracers, params)

google3/third_party/py/jax/interpreters/partial_eval.py in _remat_partial_eval(trace, f, tracers, params)
    513       out_flat = remat_call_p.bind(fun, *in_consts, **params)
    514   else:
--> 515     out_flat = remat_call_p.bind(fun, *in_consts, **params)
    516   out_pvs, jaxpr, env = aux()
    517   env = map(trace.full_raise, env)

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **params)
    606   else:
    607     tracers = map(top_trace.full_raise, args)
--> 608     outs = map(full_lower, top_trace.process_call(primitive, f, tracers, params))
    609   return apply_todos(env_trace_todo(), outs)
    610 

google3/third_party/py/jax/interpreters/partial_eval.py in process_call(self, call_primitive, f, tracers, params)
    119       tracers = map(self.instantiate_const_abstracted, tracers)
    120     if call_primitive in call_partial_eval_rules:
--> 121       return call_partial_eval_rules[call_primitive](self, f, tracers, params)
    122     if call_primitive in map_primitives:
    123       return self.process_map(call_primitive, f, tracers, params)

google3/third_party/py/jax/interpreters/partial_eval.py in _remat_partial_eval(trace, f, tracers, params)
    513       out_flat = remat_call_p.bind(fun, *in_consts, **params)
    514   else:
--> 515     out_flat = remat_call_p.bind(fun, *in_consts, **params)
    516   out_pvs, jaxpr, env = aux()
    517   env = map(trace.full_raise, env)

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **params)
    603   if top_trace is None:
    604     with new_sublevel():
--> 605       outs = primitive.impl(f, *args, **params)
    606   else:
    607     tracers = map(top_trace.full_raise, args)

google3/third_party/py/jax/core.py in call_impl(***failed resolving arguments***)
    612 def call_impl(f, *args, **params):
    613   del params  # params parameterize the call primitive, not the function
--> 614   return f.call_wrapped(*args)
    615 
    616 

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    150     gen = None
    151 
--> 152     ans = self.f(*args, **dict(self.params, **kwargs))
    153     del args
    154     while stack:

google3/third_party/py/jax/interpreters/ad.py in backward_pass(jaxpr, consts, freevar_vals, args, cotangents_in)
    220       map(write_cotangent, bound_vars, ct_free_vars_out)
    221     else:
--> 222       cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals, **eqn.params)
    223     cts_out = [zero] * len(eqn.invars) if cts_out is zero else cts_out
    224     map(write_cotangent, eqn.invars, cts_out)

google3/third_party/py/jax/lax/lax_control_flow.py in _scan_transpose(cts, *args, **kwargs)
    814     raise NotImplementedError
    815   if not all(init_lin):
--> 816     raise NotImplementedError
    817 
    818   consts, init, xs = split_list(args, [num_consts, num_carry])

NotImplementedError: 
```
Sorry for being slow on this.

Here's a crazy-sounding fix: just delete the lines that raise the exception, i.e. [these lines in lax_control_flow.py](https://github.com/google/jax/blob/efbdaf66bfa584cc635092919a23b684c7fb2247/jax/lax/lax_control_flow.py#L1007-L1008).

In JAX we model reverse-mode autodiff, like `grad`, as three steps: linearization (forward-mode autodiff), partial evaluation, and transposition (which turns forward into reverse). The exception being raised here is saying that the transpose rule can't be sure, based on the information we handed it, that the linearized-and-partial-evaluated scan we gave it is truly linear / transposable.

But anything produced by the first two steps will be transposable! We just haven't yet set it up in a way that the transpose rule can check that fact.

If you delete the exception, the transpose rule just proceeds assuming that the transposition is going to work. And it does! (If it didn't, we'd get an error. And you can always use `check_grads` from `jax.test_util` to check things numerically.) It'll always succeed for jaxprs produced by linearize-and-partial-eval. The only time it'd fail is if, hypothetically, a user tried to directly transpose their own jaxpr (not produced automatically by linearize-and-partial-eval) involving a scan which wasn't actually linear.

I think our intent with the NotImplementedError is that we wanted to add more explicit information to certify that the jaxpr to be transposed is really transposable. That still seems like a noble goal...

I want to check this thinking tomorrow, if I get a chance to work on this.
This bug appears to comes up in even the simplest examples of calling `scan` inside `remat`:
```python
import jax
from jax import lax

def scan_bug(x0):
  f = lambda x, _: (x + 1, None)

  def scanned_f(x, _):
    return lax.scan(f, x, xs=None, length=1)[0], None

  # The first three versions work. The last one crashes.
  # x, _ = f(x0, None)
  # x, _ = jax.remat(f)(x0, None)
  # x, _ = scanned_f(x0, None)
  x, _ = jax.remat(scanned_f)(x0, None)
  return x

jax.grad(scan_bug)(1.0)  # NotImplementedError
```

Resolving patch:
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -183,21 +183,12 @@ def is_linear(var):
         else:
           write_primal(eqn.outvars[0], ans)
     else:
-      call_jaxpr = eqn.params["call_jaxpr"]
+      call_jaxpr, params = core.extract_call_jaxpr(eqn.primitive, eqn.params)
       if any(is_linear(v) for v in eqn.invars):
         linear_eqns.append(eqn)
-      elif eqn.primitive is not pe.remat_call_p:
-        ans = _eval_subjaxpr_primals(
-            eqn.primitive, call_jaxpr,
-            map(read_primal, eqn.invars), eqn.params)
-        map(write_primal, eqn.outvars, ans)
-
-      # we special-case remat_call here because it can be mixed linear /
-      # nonlinear, so we always evaluate it even if it has a linear part
-      if eqn.primitive is pe.remat_call_p:
-        ans = _eval_subjaxpr_primals(
-            eqn.primitive, call_jaxpr,
-            map(read_primal, eqn.invars), 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)
 
   ct_env = {}
@@ -260,12 +251,10 @@ def is_linear(var):
         else:
           write_primal(eqn.outvars[0], ans)
     else:
-      call_jaxpr = eqn.params["call_jaxpr"]
-      if (eqn.primitive is pe.remat_call_p or
-          not any(is_linear(v) for v in eqn.invars)):
-        ans = _eval_subjaxpr_primals(
-            eqn.primitive, call_jaxpr,
-            map(read_primal, eqn.invars), eqn.params)
+      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)
 
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -1006,11 +1006,11 @@ def _scan_transpose(cts, *args, forward, length, num_consts, num_carry, jaxpr, l
   if xs_lin != [True] * (len(xs_lin) - num_eres) + [False] * num_eres:
     raise NotImplementedError
   if not all(init_lin):
-    raise NotImplementedError
+    pass  # TODO(mattjj): error check https://github.com/google/jax/issues/1963
 
-  consts, init, xs = split_list(args, [num_consts, num_carry])
-  ires, consts = split_list(consts, [num_ires])
-  xs, eres = split_list(xs, [sum(xs_lin)])
+  consts, _, xs = split_list(args, [num_consts, num_carry])
+  ires, _ = split_list(consts, [num_ires])
+  _, eres = split_list(xs, [sum(xs_lin)])
   assert not any(r is ad.undefined_primal for r in ires)
   assert not any(r is ad.undefined_primal for r in eres)
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1660,6 +1660,45 @@ def loss(u0, target, steps, dt=1/np.sqrt(2), c=1):
     u0 = np.ones_like(target)
     loss(u0, target, 10)  # doesn't crash
 
+  def test_remat_jit3(self):
+    # https://github.com/google/jax/issues/2180
+    def f(w, x):
+      a = np.dot(x, w)
+      b = np.einsum("btd,bTd->btT", a, a)
+      c = np.einsum("btT,btd->btd", b, a)
+      return np.sum(c)
+
+    w = np.ones([1, 1])
+    x = np.ones([1, 1, 1])
+    f = api.remat(f)
+    api.grad(f)(w, x)  # doesn't crash
+
+    @api.jit
+    def mul(a, b):
+      return a * b
+
+    def f(w, x):
+      a = mul(w, x)
+      b = mul(a, a)
+      return b
+
+    w = 1.
+    x = 1.
+    f = api.remat(f)
+    api.grad(f)(w, x)  # doesn't crash
+
+  def test_remat_scan2(self):
+    # https://github.com/google/jax/issues/1963
+
+    def scan_bug(x0):
+      f = lambda x, _: (x + 1, None)
+      def scanned_f(x, _):
+        return lax.scan(f, x, xs=None, length=1)[0], None
+      x, _ = jax.remat(scanned_f)(x0, None)
+      return x
+
+    jax.grad(scan_bug)(1.0)  # doesn't crash
+
   def test_trivial_computations(self):
     x = np.array([1, 2, 3])
     y = api.jit(lambda x: x)(x)