Repository: google/jax

Issue:
remat inside scan only works when computing gradient
I tried using `remat` inside a variation of my TPU wave equation example (for optimizing an initial condition). It works when computing a gradient, but not when doing the forwards calculation of the loss:
```python
from functools import partial
import jax
from jax import lax
import jax.numpy as jnp
import numpy as np
from scipy import ndimage

def axis_slice(ndim, index, axis):
  slices = [slice(None)] * ndim
  slices[axis] = index
  return tuple(slices)

def slice_along_axis(array, index, axis):
  return array[axis_slice(array.ndim, index, axis)]

def shift(array, offset, axis):
  index = slice(offset, None) if offset >= 0 else slice(None, offset)
  sliced = slice_along_axis(array, index, axis)
  padding = [(0, 0)] * array.ndim
  padding[axis] = (-min(offset, 0), max(offset, 0))
  return jnp.pad(sliced, padding, mode='constant', constant_values=0)

def laplacian(array, step=1):
  left = shift(array, +1, axis=0)
  right = shift(array, -1, axis=0)
  up = shift(array, +1, axis=1)
  down = shift(array, -1, axis=1)
  convolved = (left + right + up + down - 4 * array)
  if step != 1:
    convolved *= (1 / step ** 2)
  return convolved

def scalar_wave_equation(u, c=1, dx=1):
  return c ** 2 * laplacian(u, dx)

@jax.remat
def leapfrog_step(state, dt=0.5, c=1):
  # https://en.wikipedia.org/wiki/Leapfrog_integration
  u, u_t = state
  u_tt = scalar_wave_equation(u, c)
  u_t = u_t + u_tt * dt
  u = u + u_t * dt
  return (u, u_t)

@partial(jax.jit, static_argnums=(1,))
def _multi_step(state, count, dt, c):
  f = lambda s, _: (leapfrog_step(s, dt, c), _)
  return lax.scan(f, state, None, count)

def multi_step(state, count, dt=1/np.sqrt(2), c=1):
  return _multi_step(state, count, dt, c)

def loss(u0, target, steps, dt=1/np.sqrt(2), c=1):
  init = (u0, jnp.zeros_like(u0))
  (uf, _), _ = multi_step(init, steps, dt, c)
  return ((uf - target) ** 2).mean()

grad_loss = jax.jit(jax.grad(loss), static_argnums=(2,))

x = np.linspace(0, 1, num=128, endpoint=False)
y = np.linspace(0, 1, num=128, endpoint=False)
x_mesh, y_mesh = np.meshgrid(x, y, indexing='ij')

target = ndimage.gaussian_filter(
    1.0 * (((x_mesh - 1/3) ** 2 + (y_mesh - 1/4) ** 2) < 0.1 ** 2),
    sigma=1)
u0 = np.ones_like(target)

grad_loss(u0, target, 10)  # works
jax.value_and_grad(loss)(u0, target, 10)  # also works
loss(u0, target, 10)  # doesn't work
```
Here's the traceback:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-9-81418533a6e4> in <module>
----> 1 loss(u0, target, 10)  # doesn't work

<ipython-input-8-b038831f9764> in loss(u0, target, steps, dt, c)
     54 def loss(u0, target, steps, dt=1/np.sqrt(2), c=1):
     55   init = (u0, jnp.zeros_like(u0))
---> 56   (uf, _), _ = multi_step(init, steps, dt, c)
     57   return ((uf - target) ** 2).mean()
     58

<ipython-input-8-b038831f9764> in multi_step(state, count, dt, c)
     50
     51 def multi_step(state, count, dt=1/np.sqrt(2), c=1):
---> 52   return _multi_step(state, count, dt, c)
     53
     54 def loss(u0, target, steps, dt=1/np.sqrt(2), c=1):

~/dev/jax/jax/api.py in f_jitted(*args, **kwargs)
    147     _check_args(args_flat)
    148     flat_fun, out_tree = flatten_fun(f, in_tree)
--> 149     out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend)
    150     return tree_unflatten(out_tree(), out)
    151

~/dev/jax/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)

~/dev/jax/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)

~/dev/jax/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

~/dev/jax/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

~/dev/jax/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:

<ipython-input-8-b038831f9764> in _multi_step(state, count, dt, c)
     47 def _multi_step(state, count, dt, c):
     48   f = lambda s, _: (leapfrog_step(s, dt, c), _)
---> 49   return lax.scan(f, state, None, count)
     50
     51 def multi_step(state, count, dt=1/np.sqrt(2), c=1):

~/dev/jax/jax/lax/lax_control_flow.py in scan(f, init, xs, length)
    612                     forward=True, length=length, jaxpr=jaxpr,
    613                     num_consts=len(consts), num_carry=len(init_flat),
--> 614                     linear=(False,) * (len(consts) + len(in_flat)))
    615   return tree_unflatten(out_tree, out)
    616

~/dev/jax/jax/lax/lax_control_flow.py in scan_bind(*args, **kwargs)
    979   return core.Primitive.bind(scan_p, *args, forward=forward, length=length,
    980                              jaxpr=jaxpr, num_consts=num_consts,
--> 981                              num_carry=num_carry, linear=linear)
    982
    983 scan_p = core.Primitive("scan")

~/dev/jax/jax/core.py in bind(self, *args, **kwargs)
    160
    161     tracers = map(top_trace.full_raise, args)
--> 162     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    163     if self.multiple_results:
    164       return map(full_lower, out_tracer)

~/dev/jax/jax/interpreters/partial_eval.py in process_primitive(self, primitive, tracers, params)
     96   def process_primitive(self, primitive, tracers, params):
     97     if primitive in custom_partial_eval_rules:
---> 98       return custom_partial_eval_rules[primitive](self, *tracers, **params)
     99     else:
    100       pvs, consts = unzip2(t.pval for t in tracers)

~/dev/jax/jax/lax/lax_control_flow.py in _scan_partial_eval(trace, *tracers, **kwargs)
    729     unknowns = const_uk + carry_uk + xs_uk
    730     jaxpr_1, jaxpr_2, out_uk = pe.partial_eval_jaxpr(
--> 731         jaxpr, unknowns, instantiate=carry_uk + [False] * num_ys)
    732     carry_uk_out, ys_uk = out_uk[:num_carry], out_uk[num_carry:]
    733     if carry_uk_out == carry_uk:

~/dev/jax/jax/interpreters/partial_eval.py in partial_eval_jaxpr(jaxpr, unknowns, instantiate)
    453   pvals = [PartialVal((abstract_unit, unit)) if uk else PartialVal((aval, unit))
    454            for aval, uk in zip(jaxpr.in_avals, unknowns)]
--> 455   jaxpr_1, out_pvals, consts_1 = trace_to_jaxpr(lu.wrap_init(fun), pvals, instantiate=True)
    456   (out_pvs_2, jaxpr_2, num_res), = cell
    457   assert len(jaxpr_2.constvars) == num_res

~/dev/jax/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

~/dev/jax/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:

~/dev/jax/jax/interpreters/partial_eval.py in fun(*vals)
    446     pvals = [PartialVal((aval, unit)) if uk else PartialVal((None, val))
    447              for aval, val, uk in zip(jaxpr.in_avals, vals, unknowns)]
--> 448     jaxpr_2, out_pvals_2, consts_2 = trace_to_jaxpr(f, pvals, instantiate=instantiate)
    449     out_pvs_2, out_consts_2 = unzip2(out_pvals_2)
    450     cell.append((out_pvs_2, jaxpr_2, len(consts_2)))

~/dev/jax/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

~/dev/jax/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:

~/dev/jax/jax/core.py in jaxpr_as_fun(typed_jaxpr, *args)
     77 @curry
     78 def jaxpr_as_fun(typed_jaxpr, *args):
---> 79   return eval_jaxpr(typed_jaxpr.jaxpr, typed_jaxpr.literals, (), *args)
     80
     81

~/dev/jax/jax/core.py in eval_jaxpr(jaxpr, consts, freevar_vals, *args)
    212                in eqn.bound_subjaxprs]
    213     subfuns = map(lu.wrap_init, subfuns)
--> 214     ans = eqn.primitive.bind(*(subfuns + in_vals), **eqn.params)
    215     if eqn.primitive.multiple_results:
    216       map(write, eqn.outvars, ans)

~/dev/jax/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

~/dev/jax/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)

~/dev/jax/jax/interpreters/partial_eval.py in _remat_partial_eval(trace, f, tracers, params)
    522   # outputs are known/unknown, we use partial_eval_jaxpr to get out_unknowns.
    523   jaxpr_converted = convert_freevars_jaxpr(jaxpr)
--> 524   in_avals = ([raise_to_shaped(t.pval[0]) for t in env]
    525               + [raise_to_shaped(pv) for pv in in_pvs])
    526   out_avals = [raise_to_shaped(pv if pv is not None

~/dev/jax/jax/interpreters/partial_eval.py in <listcomp>(.0)
    522   # outputs are known/unknown, we use partial_eval_jaxpr to get out_unknowns.
    523   jaxpr_converted = convert_freevars_jaxpr(jaxpr)
--> 524   in_avals = ([raise_to_shaped(t.pval[0]) for t in env]
    525               + [raise_to_shaped(pv) for pv in in_pvs])
    526   out_avals = [raise_to_shaped(pv if pv is not None

~/dev/jax/jax/abstract_arrays.py in raise_to_shaped(aval, weak_type)
    232     return abstract_token
    233   else:
--> 234     raise TypeError(type(aval))
    235
    236 core.literalable_types.update(array_types)

TypeError: <class 'NoneType'>
```


Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -2027,7 +2027,8 @@ def checkpoint(fun, concrete=False):
   def fun_remat(*args, **kwargs):
     args_flat, in_tree = tree_flatten((args, kwargs))
     flat_fun, out_tree = flatten_fun(lu.wrap_init(fun), in_tree)
-    out_flat = pe.remat_call(flat_fun, *args_flat, concrete=concrete)
+    out_flat = pe.remat_call(flat_fun, *args_flat, name=flat_fun.__name__,
+                             concrete=concrete)
     return tree_unflatten(out_tree(), out_flat)
   return fun_remat
 remat = checkpoint
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -526,7 +526,7 @@ def _remat_partial_eval(trace, f, tracers, params):
   # Since we traced with everything marked as unknown, but we need to know which
   # outputs are known/unknown, we use partial_eval_jaxpr to get out_unknowns.
   jaxpr_converted = convert_freevars_jaxpr(jaxpr)
-  in_avals = ([raise_to_shaped(t.pval[0]) for t in env]
+  in_avals = ([raise_to_shaped(partial_val_aval(*t.pval)) for t in env]
               + [raise_to_shaped(pv) for pv in in_pvs])
   out_avals = [raise_to_shaped(pv if pv is not None
                                else abstract_unit if var is unitvar


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1600,6 +1600,46 @@ def move(R,i):
 
     api.grad(func)(5.0)  # doesn't crash
 
+  def test_remat_jit2(self):
+    @api.jit
+    def f(x):
+      y = 2 * x
+
+      @api.remat
+      def g():
+        return y
+
+      return g()
+
+    self.assertAllClose(f(3), 6, check_dtypes=False)
+
+  def test_remat_nontrivial_env(self):
+    # simplified from https://github.com/google/jax/issues/2030
+
+    @api.remat
+    def foo(state, dt=0.5, c=1):
+      u, u_t = state
+      u_tt = c**2 * u
+      u_t = u_t + u_tt * dt
+      return (u, u_t)
+
+    @partial(api.jit, static_argnums=(1,))
+    def _multi_step(state, count, dt, c):
+      f = lambda s, _: (foo(s, dt, c), _)
+      return lax.scan(f, state, None, count)
+
+    def multi_step(state, count, dt=1/np.sqrt(2), c=1):
+      return _multi_step(state, count, dt, c)
+
+    def loss(u0, target, steps, dt=1/np.sqrt(2), c=1):
+      init = (u0, np.zeros_like(u0))
+      (uf, _), _ = multi_step(init, steps, dt, c)
+      return ((uf - target) ** 2).mean()
+
+    target = np.zeros((128, 128))
+    u0 = np.ones_like(target)
+    loss(u0, target, 10)  # doesn't crash
+
   def test_trivial_computations(self):
     x = np.array([1, 2, 3])
     y = api.jit(lambda x: x)(x)