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]
```


Discussion:
I'm seeing this happen in yet another test case. I'll keep looking.
Thanks for raising this!

Minimal test cases are easier for us to make progress on. Any chance you can pare that down?
Sorry for being lazy/sloppy - I'll make more self contained repros in the future. I've updated the original issue with a much smaller test case. I think issue is that somewhere along the way the iteration count is being casted to a vanilla scalar (as opposed to a zero sized jax IntArray[]).

Resolving patch:
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- 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


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -592,6 +592,37 @@ def loss(params, inputs, targets):
     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()