Repository: google/jax
Issue: more scan jacrev/jacfwd bugs

This is a follow up of #711 (different error from #757):
``` 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)

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.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, (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 - get_params(opt_final)[0]

loss_opt_init, loss_opt_update, loss_get_params = optimizers.sgd(5e-2)

jax.grad(loss, argnums=(0,))(initial_params) # okay
jax.jacrev(loss, argnums=(0,))(initial_params) # AttributeError: 'JaxTuple' object has no attribute 'broadcast'
jax.jacfwd(loss, argnums=(0,))(initial_params) # TypeError: <class 'tuple'>
```

Changed files:
jax/interpreters/batching.py, tests/optimizers_test.py

diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
index d795c56f3..cb1d9ccc5 100644
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -252,14 +252,16 @@ def reducer_batcher(prim, batched_args, batch_dims, axes, **params):
 
 def add_batched(batched_args, batch_dims):
   bdx, bdy = batch_dims
-  xs, ys = batched_args
-  if bdx == bdy:
-    return add_jaxvals_p.bind(xs, ys), bdx
+  x, y = batched_args
+  x_aval, y_aval = map(get_aval, batched_args)
+  assert core.skip_checks or x_aval == y_aval
+  if bdx == bdy or x_aval == AbstractTuple(()):
+    return add_jaxvals_p.bind(x, y), bdx
   else:
-    sz = (dimsize(bdx, xs) | dimsize(bdy, ys)).pop()
+    sz = (dimsize(bdx, x) | dimsize(bdy, y)).pop()
     move_bdim = partial(bdim_at_front, broadcast_size=sz, force_broadcast=True)
-    xs, ys = map(move_bdim, batched_args, batch_dims)
-    return add_jaxvals_p.bind(xs, ys), 0
+    x, y = map(move_bdim, batched_args, batch_dims)
+    return add_jaxvals_p.bind(x, y), 0
 primitive_batchers[add_jaxvals_p] = add_batched
 
 def zeros_like_batched(batched_args, batch_dims):
@@ -409,23 +411,22 @@ def instantiate_bdim(size, axis, instantiate, bdim, x):
   """
   def _inst(instantiate, bdim, x):
     if type(instantiate) is tuple:
-      assert type(bdim) is tuple  # instantiate at same granularity as bdim
-      return core.pack(map(_inst, instantiate, bdim, x))
+      if type(bdim) is tuple:
+        return core.pack(map(_inst, instantiate, bdim, x))
+      elif type(bdim) is int or bdim is None:
+        bdims = (bdim,) * len(instantiate)
+        return core.pack(map(_inst, instantiate, bdims, x))
+      else:
+        raise TypeError(type(bdim))
     elif type(instantiate) is bool:
-      if not instantiate:
-        if bdim is None:
-          return x
-        elif type(bdim) is int:
-          return moveaxis2(bdim, axis, x)
-        else:
-          raise TypeError(type(bdim))
+      if bdim is None:
+        return broadcast2(size, axis, x) if instantiate else x
+      elif type(bdim) is int:
+        return moveaxis2(bdim, axis, x)
+      elif type(bdim) is tuple:
+        return pack(map(partial(_inst, instantiate), bdim, x))
       else:
-        if bdim is None:
-          return broadcast2(size, axis, x)
-        elif type(bdim) is int:
-          return moveaxis2(bdim, axis, x)
-        else:
-          raise TypeError(type(bdim))
+        raise TypeError(type(bdim))
     else:
       raise TypeError(type(instantiate))
 
@@ -449,7 +450,7 @@ def broadcast2(size, axis, x):
   return _broadcast2(size, axis, x, get_aval(x))
 
 def _broadcast2(size, axis, x, aval):
-  if type(aval) is JaxTuple:
+  if type(aval) is AbstractTuple:
     return core.pack(map(partial(_broadcast2, size, axis), x, aval))
   else:
     # see comment at the top of this section
@@ -460,13 +461,13 @@ def _broadcast2(size, axis, x, aval):
 
 def _promote_aval_rank(n, batched, aval):
   assert isinstance(aval, core.AbstractValue)
-  if isinstance(aval, core.AbstractTuple):
+  if isinstance(aval, AbstractTuple):
     t = type(batched)
     if t is tuple:
-      return core.AbstractTuple(map(partial(_promote_aval_rank, n), batched, aval))
+      return AbstractTuple(map(partial(_promote_aval_rank, n), batched, aval))
     elif t is bool:
       if batched:
-        return core.AbstractTuple(map(partial(_promote_aval_rank, n, batched), aval))
+        return AbstractTuple(map(partial(_promote_aval_rank, n, batched), aval))
       else:
         return aval
     else:
diff --git a/tests/optimizers_test.py b/tests/optimizers_test.py
index d3995160e..e8725b94d 100644
--- a/tests/optimizers_test.py
+++ b/tests/optimizers_test.py
@@ -24,8 +24,9 @@ import numpy as onp
 
 import jax.numpy as np
 import jax.test_util as jtu
-from jax import jit, grad
+from jax import jit, grad, jacfwd, jacrev
 from jax import core, tree_util
+from jax import lax
 from jax.experimental import optimizers
 from jax.interpreters import xla
 
@@ -244,7 +245,48 @@ class OptimizerTests(jtu.JaxTestCase):
     expected = 0.9 * norm
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testIssue758(self):
+    # code from https://github.com/google/jax/issues/758
+    # this is more of a scan + jacfwd/jacrev test, but it lives here to use the
+    # optimizers.py code
 
+    def harmonic_bond(conf, params):
+      return np.sum(conf * params)
+
+    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 = grad(energy_fn, argnums=(0,))
+      opt_state = opt_init(x0)
+
+      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, _ = lax.scan(apply_carry, (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 - get_params(opt_final)[0]
+
+    loss_opt_init, loss_opt_update, loss_get_params = optimizers.sgd(5e-2)
+
+    J1 = jacrev(loss, argnums=(0,))(initial_params)
+    J2 = jacfwd(loss, argnums=(0,))(initial_params)
+    self.assertAllClose(J1, J2, check_dtypes=True)
 
 
 if __name__ == '__main__':