Repository: google/jax

Issue:
Meta-Gradient causes "TypeError: <class 'jax.ad_util.Zero'> is not a valid JAX type"
Hello,

I discovered the following unexpected behaviour while calculating meta-gradients (gradients through gradients) with JAX.

```
import jax
from jax import lax
from jax import numpy as jnp

def loss(params, meta_params):
  (net, mpo) = params
  return meta_params * mpo * net

def inner(meta_params, params):
  grads = jax.grad(loss)(params, meta_params)
  grads = lax.psum(grads, axis_name="i")
  net_grads, mpo_grads = grads
  net = params[0] + net_grads
  mpo = params[1]                           # Does not work!
  # mpo = params[1] + mpo_grads  # Works if I add mpo_grads
  return mpo * net

def outer(params):
  meta_params = jnp.array(1.0)
  return jax.grad(inner)(meta_params, params)

params = (jnp.array([1.0]), jnp.array([1.0]))
learner_output = jax.pmap(outer, axis_name='i')(params)
```

Outputs

```
[....]
jax/core.py in concrete_aval(x)
    778     handler = pytype_aval_mappings.get(typ)
    779     if handler: return handler(x)
--> 780   raise TypeError(f"{type(x)} is not a valid JAX type")
    781 
    782 

TypeError: <class 'jax.ad_util.Zero'> is not a valid JAX type
```

This only appears inside a pmap (both CPU and TPU as backend throw the same error).
The same code works without pmap, when using jit (or vmap), or if the psum is removed.


Discussion:
A possible workaround seems to be to first split the gradients and then aggregate them:

```
net_grads, mpo_grads = grads
net_grads = lax.psum(net_grads, axis_name="i")
mpo_grads = lax.psum(mpo_grads, axis_name="i")
```

Thanks for the report, and nice repro! This looks like one of our symbolic zeros leaking somewhere that it shouldn't be.

Resolving patch:
diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
--- a/jax/lax/lax_parallel.py
+++ b/jax/lax/lax_parallel.py
@@ -364,6 +364,12 @@ def _translate(val):
       return psum(val)
   return xops.Tuple(c, list(map(_translate, args)))
 
+def _psum_transpose_rule(cts, axis_name, axis_index_groups):
+  nonzero_out_cts, treedef = tree_util.tree_flatten(cts)
+  nonzero_in_cts = psum_p.bind(*nonzero_out_cts, axis_name=axis_name,
+                               axis_index_groups=axis_index_groups)
+  return tree_util.tree_unflatten(treedef, nonzero_in_cts)
+
 psum_p = standard_pmap_primitive('psum', multiple_results=True)
 psum_p.def_abstract_eval(
     lambda *args, **params: tuple(map(raise_to_shaped, args)))
@@ -371,8 +377,7 @@ def _translate(val):
     partial(_allreduce_split_axis_rule, psum_p, lax._reduce_sum)
 xla.parallel_translations[psum_p] = _psum_translation_rule
 pxla.parallel_pure_rules[psum_p] = lambda *args, shape: (x * prod(shape) for x in args)
-ad.deflinear(psum_p, lambda ts, axis_name, axis_index_groups: psum_p.bind(
-    *ts, axis_name=axis_name, axis_index_groups=axis_index_groups))
+ad.deflinear(psum_p, _psum_transpose_rule)
 pxla.multi_host_supported_collectives.add(psum_p)
 
 


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -1290,6 +1290,30 @@ def foo(x): return x
       self.assertIn("The jitted function foo includes a pmap",
                     str(w[-1].message))
 
+  def testPsumZeroCotangents(self):
+    # https://github.com/google/jax/issues/3651
+    def loss(params, meta_params):
+      (net, mpo) = params
+      return meta_params * mpo * net
+
+    def inner(meta_params, params):
+      grads = jax.grad(loss)(params, meta_params)
+      grads = lax.psum(grads, axis_name="i")
+      net_grads, mpo_grads = grads
+      net = params[0] + net_grads
+      mpo = params[1]
+      return mpo * net
+
+    def outer(params):
+      meta_params = jnp.array(4.0)
+      return jax.grad(inner)(meta_params, params)
+
+    params = (jnp.array([2.0]), jnp.array([3.0]))
+    jax.pmap(outer, axis_name='i')(params)  # doesn't crash
+
+    f = jax.pmap(outer, axis_name='i')
+    jtu.check_grads(f, (params,), 2, ["fwd", "rev"], 1e-3, 1e-3)
+
 
 class VmapOfPmapTest(jtu.JaxTestCase):