Repository: google/jax
Issue: lax_scan/lattice_join shape inconsistencies

For certain native scalar types, the newly implemented lattice_join code doesn't seem to be able to match the types. This is a follow-up of #650 

``` python
    def lattice_join(x, y):
      if x is None:
        return y
      elif y is None:
        return x
      elif isinstance(x, type(y)):
        return y.join(x)
      elif isinstance(y, type(x)):
        return x.join(y)
      else:
>       raise TypeError((x, y))
E       TypeError: (ShapedArray(int64[]), ())
```

Code to reproduce:

``` python
import unittest
import numpy as onp
import jax.numpy as np
import functools

import jax
from jax.config import config; config.update("jax_enable_x64", True)
from jax.experimental import optimizers
from jax.test_util import check_grads

def harmonic_bond(conf, params):
    return np.sum(conf * params)

class TestOptimizeGeometry(unittest.TestCase):

    def test_case(self):

        opt_init, opt_update, get_params = optimizers.sgd(5e-2)

        x0 = onp.array([0.5], dtype=onp.float64)
        params = onp.array([0.3], dtype=onp.float64)

        def minimize_structure(test_params):
            energy_fn = functools.partial(harmonic_bond, params=test_params)
            grad_fn = jax.jit(jax.grad(energy_fn, argnums=(0,)))
            opt_state = opt_init(x0)

            # use lax.scan, way faster compilation times.
            def apply_carry(carry, _):
                i, x = carry
                g = grad_fn(get_params(x))[0]
                new_state = opt_update(i, g, x)
                new_carry = (i+1, new_state)
                return new_carry, _

            carry_final, _ = jax.lax.scan(apply_carry, (np.array(0), opt_state), np.zeros((75, 0)))

            trip, opt_final = carry_final

            assert trip == 75

            return opt_final

        initial_params = 0.5
        minimize_structure(initial_params)

        def loss(test_params):
            opt_final = minimize_structure(test_params)
            return 1.0-opt_final

        loss_opt_init, loss_opt_update, loss_get_params = optimizers.sgd(5e-2)

        loss_grad_fn = jax.grad(loss, argnums=(0,))
        loss_opt_state = loss_opt_init(initial_params)
        loss_params = loss_get_params(loss_opt_state)
        loss_grad = loss_grad_fn(loss_params)[0]
```

Changed files:
jax/lax/lax_control_flow.py, tests/lax_control_flow_test.py

diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
index f29630f2b..b69979173 100644
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -428,16 +428,25 @@ def _maybe_tracer_tuple_to_abstract_tuple(tup):
 
 ### scan
 
-def _convert_zeros(convert_symbolic, example, tangent):
-  if tangent is ad.zero:
-    if not convert_symbolic:
+def _convert_zeros(instantiate, example, tangent):
+  t = type(instantiate)
+  if t is bool:
+    if instantiate:
+      return ad.instantiate_zeros(example, tangent)
+    elif tangent is ad_util.zero:
       return core.unit
     else:
-      return ad.zeros_like_jaxval(example)
-  elif type(tangent) is ad.TangentTuple:
-    return core.pack(map(_convert_zeros, convert_symbolic, example, tangent))
+      raise TypeError(tangent)  # not clear if ever reachable
+  elif t is tuple:
+    if type(tangent) is ad.TangentTuple:
+      return core.pack(map(_convert_zeros, instantiate, example, tangent))
+    elif tangent is ad_util.zero:
+      zeros = [ad_util.zero] * len(instantiate)
+      return core.pack(map(_convert_zeros, instantiate, example, zeros))
+    else:
+      raise TypeError(tangent)
   else:
-    return tangent
+    raise TypeError(t)
 
 def _demote_aval_rank(xs):
   assert isinstance(xs, core.AbstractValue)
@@ -641,7 +650,7 @@ def _scan_partial_eval(trace, *tracers, **kwargs):
   length = kwargs.pop('length')
   forward = kwargs.pop('forward')
   assert not kwargs
-  in_pvs, in_consts = unzip2([t.pval for t in tracers])
+  in_pvs, _ = unzip2([t.pval for t in tracers])
   sc_consts, sc_init, sc_xs = map(pe.unknown, in_pvs)
 
   sc_carry = sc_init
@@ -819,7 +828,19 @@ def _make_typed_jaxpr(traceable, in_avals):
 class FixedPointError(Exception): pass
 
 
+# We use a custom bind for scan just to add some error checks
+def scan_bind(consts, init, xs, forward, length, jaxpr):
+  if not core.skip_checks:
+    assert type(jaxpr.in_avals) is tuple
+    consts_aval, init_aval, xs_aval = jaxpr.in_avals
+    assert type(jaxpr.out_aval) is core.AbstractTuple
+    carry_aval, y_aval = jaxpr.out_aval
+    assert init_aval == carry_aval
+  return core.Primitive.bind(scan_p, consts, init, xs,
+                             forward=forward, length=length, jaxpr=jaxpr)
+
 scan_p = core.Primitive("scan")
+scan_p.def_custom_bind(scan_bind)
 scan_p.def_impl(_scan_impl)
 ad.primitive_jvps[scan_p] = _scan_jvp
 ad.primitive_transposes[scan_p] = _scan_transpose
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index 6e7b4848b..9d1def88e 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -592,6 +592,37 @@ class LaxControlFlowTest(jtu.JaxTestCase):
     expected = (onp.zeros_like(W_trans), onp.zeros_like(W_out))
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testIssue711(self):
+    # Tests reverse-mode differentiation through a scan for which the scanned
+    # function also involves reverse-mode differentiation.
+    # See https://github.com/google/jax/issues/711
+    def harmonic_bond(conf, params):
+      return np.sum(conf * params)
+
+    def minimize_structure(test_params):
+      energy_fn = partial(harmonic_bond, params=test_params)
+      grad_fn = api.grad(energy_fn)
+
+      def apply_carry(carry, _):
+        i, x = carry
+        new_x = x - 0.1 * api.grad(energy_fn)(x)
+        new_carry = (i+1, new_x)
+        return new_carry, _
+
+      x0 = np.array([1., 2., 3.])
+      carry_final, _ = lax.scan(apply_carry, (0, x0), np.zeros((75, 0)))
+      _, x_final = carry_final
+      return x_final
+
+    initial_params = 0.5
+    minimize_structure(initial_params)  # doesn't crash
+
+    def loss(test_params):
+      x_final = minimize_structure(test_params)
+      return np.sum(np.sin(1.0 - x_final))
+
+    api.grad(loss)(0.25)  # doesn't crash
+
 
 if __name__ == '__main__':
   absltest.main()