Repository: google/jax

Issue:
remat with scan
First, let me express my excitement with the new remat functionality for automatic checkpointing! Really awesome :)

As I begin testing this out for my purposes, I am running into an issue that pops up when using lax.scan, though I'm not sure exactly where its coming from. I created a minimal example (see below), which I have tested on Colab. When both of the flags `test_remat` and `test_scan` are set to `True`, an error ultimately saying `TypeError: <class 'jax.ad_util.Zero'> is not a valid Jax type` appears. I should note that on a (much) more complicated version of this I was getting a NotImplemented error that I think is related, but I wasn't able to reproduce that in a minimal code so its possible its something else. I'd be happy to share if that would help. 

Any ideas on what is going on? Thanks in advance!


Here's the minimal example:

```
!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, grad, lax, remat

########################################
#Error only when both of these are True#
########################################
test_remat = True
test_scan = True

key = random.PRNGKey(0)
key, split = random.split(key) 
n=5

def func(D0):
  def shift(R, dR, **unused_kwargs):
    return R + dR

  def apply_fn(R):
    return D0*R

  Rinit = random.uniform(split, (n,3), minval=0.0, maxval=5.0, dtype=np.float32)

  def move(R,i):
    F = apply_fn(R)
    return shift(R, 0.001*F), np.array([0.])
  
  if(test_remat):
    move = remat(move)
  
  if(test_scan):
    R, temp = lax.scan(move, Rinit, np.arange(2))
  else:
    R, temp = move(move(Rinit,0)[0],0)

  return R[0,0]


grad(func)(5.0)
```


Discussion:
Thanks for the enthusiasm, and the clear repro!

This error pops up when we aren't handling the AD system's symbolic zeros correctly. Usually the fix is trivial once we spot it...

Resolving patch:
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
@@ -524,8 +524,11 @@ def _remat_partial_eval(trace, f, tracers, params):
   jaxpr_converted = convert_freevars_jaxpr(jaxpr)
   in_avals = ([raise_to_shaped(t.pval[0]) 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 core.get_aval(const))
-               for pv, const in zip(out_pvs, out_pval_consts1)]
+  out_avals = [raise_to_shaped(pv if pv is not None
+                               else abstract_unit if var is unitvar
+                               else get_aval(var.val) if type(var) is Literal
+                               else get_aval(const))
+               for var, pv, const in zip(jaxpr.outvars, out_pvs, out_pval_consts1)]
   typed_jaxpr = core.TypedJaxpr(jaxpr_converted, consts, in_avals, out_avals)
   in_unknowns = [t.pval[0] is not None for t in it.chain(env, tracers)]
   jaxpr_1, jaxpr_2, out_unknowns = partial_eval_jaxpr(typed_jaxpr, in_unknowns, False)
@@ -565,7 +568,7 @@ def _dce_jaxpr(typed_jaxpr, outputs):
   # TODO(mattjj): better DCE
   jaxpr = typed_jaxpr.jaxpr
   outvars, out_avals = jaxpr.outvars, typed_jaxpr.out_avals
-  out_pairs = [(var, aval) if output else (core.unitvar, core.abstract_unit)
+  out_pairs = [(var, aval) if output else (unitvar, core.abstract_unit)
               for var, aval, output in zip(outvars, out_avals, outputs)]
   new_outvars, new_out_avals = unzip2(out_pairs)
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1533,6 +1533,35 @@ def binom_checkpoint(funs):
     self.assertAllClose(f1(x), f2(x), check_dtypes=False)
     self.assertAllClose(api.grad(f1)(x), api.grad(f2)(x), check_dtypes=False)
 
+  def test_remat_symbolic_zeros(self):
+    # code from https://github.com/google/jax/issues/1907
+    test_remat = True
+    test_scan = True
+
+    key = jax.random.PRNGKey(0)
+    key, split = jax.random.split(key)
+    n = 5
+
+    def func(D0):
+      def shift(R, dR, **unused_kwargs):
+        return R + dR
+
+      def apply_fn(R):
+        return D0 * R
+
+      Rinit = jax.random.uniform(split, (n,3), minval=0.0, maxval=5.0,
+                                 dtype=np.float32)
+
+      def move(R,i):
+        F = apply_fn(R)
+        return shift(R, 0.001 * F), np.array([0.])
+
+      move = api.remat(move)
+      R, temp = lax.scan(move, Rinit, np.arange(2))
+      return R[0, 0]
+
+    api.grad(func)(5.0)  # doesn't crash
+
   def test_trivial_computations(self):
     x = np.array([1, 2, 3])
     y = api.jit(lambda x: x)(x)