Repository: google/jax

Issue:
jvp of np.sum fails in shape assertion when jit'ed
The following piece of code, which computes the jvp of the `(reduce_)sum` operation, fails with `AssertionError: (3, 2, 2) != (3,)` when the sum is jit'ed. It computes without problems and gives the expected result when the `sum` operation is not jit'ed. The assertion is thrown from the `_reduce_sum_shape_rule` (lax.py, line 3017) which appears to be only called during actual evaluation of the `f_jvp` function when jit'ed.

``` import jax
import jax.numpy as np

T = np.array([[1., 2.], [3., 4.], [5., 6.]])
x = np.array([1, 2, 3])

mysum = jax.jit(np.sum)

# this doesn't work when jit'ed
y, f_jvp = jax.linearize(mysum, x)
Z = f_jvp(T)

# this doesn't work either when jit'ed
# y, Z = jax.jvp(mysum, (x, ), (T, ))

print(y)
print(Z)
```

Unfortunately, I haven't been able to work out a fix for this on my own so far due to the many layers of abstraction and indirections in between and would be grateful if someone could push me the right way on this.


Discussion:
Thanks for raising this!

Actually, `x` and `T` should have the same shape here: in general, tangent vectors need to have the same shape as the corresponding input (primal) values. However, we're not raising this error in a good way, and it's leading to behavior that's only being caught way downstream.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -762,13 +762,23 @@ def linearize(fun, *primals):
   out_primal, out_pval, jaxpr, consts = ad.linearize(jaxtree_fun, *primals_flat)
   out_tree = out_tree()
   out_primal_py = build_tree(out_tree, out_primal)
-  lifted_jvp = partial(lift_linearized, jaxpr, consts, (in_trees, out_tree), out_pval)
+  primal_avals = list(map(core.get_aval, primals_flat))
+  lifted_jvp = partial(lift_linearized, jaxpr, primal_avals, consts,
+                       (in_trees, out_tree), out_pval)
   return out_primal_py, lifted_jvp
 
-def lift_linearized(jaxpr, consts, io_tree, out_pval, *py_args):
-  def fun(*args):
-    primals = pack(args) # doesn't matter what these are-they'll be ignored
-    tangents = pack(args)
+def lift_linearized(jaxpr, primal_avals, consts, io_tree, out_pval, *py_args):
+  def fun(*tangents):
+    tangent_avals = list(map(core.get_aval, tangents))
+    for primal_aval, tangent_aval in zip(primal_avals, tangent_avals):
+      try:
+        core.lattice_join(primal_aval, tangent_aval)
+      except TypeError:
+        msg = ("linearized function called on tangent values inconsistent with "
+               "the original primal values.")
+        raise ValueError(msg)
+    primals = pack(tangents)  # doesn't matter what these are-they'll be ignored
+    tangents = pack(tangents)
     _, ans = eval_jaxpr(jaxpr, consts, (), primals, tangents)
     return pe.merge_pvals(ans, out_pval)
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -784,6 +784,20 @@ def fun(x):
 
     assert len(api.make_jaxpr(fun)(1).eqns) == 0
 
+  def test_issue_871(self):
+    T = np.array([[1., 2.], [3., 4.], [5., 6.]])
+    x = np.array([1, 2, 3])
+
+    y, f_jvp = api.linearize(np.sum, x)
+    jtu.check_raises(lambda: f_jvp(T), ValueError,
+                     ("linearized function called on tangent values "
+                      "inconsistent with the original primal values."))
+
+    y, f_jvp = api.linearize(api.jit(np.sum), x)
+    jtu.check_raises(lambda: f_jvp(T), ValueError,
+                     ("linearized function called on tangent values "
+                      "inconsistent with the original primal values."))
+
 
 if __name__ == '__main__':
   absltest.main()