Repository: google/jax

Issue:
Scatter transpose rule doesn’t handle symbolic zeros
Unfortunately, I don't have a terribly concise repro for this error. From playing around with it, it seems as though the error probably involves index_update, but it's hard to be sure. However, given the function

```python
def nvt_nose_hoover(energy_fn, dt, T, chain_length=5, tau=0.01):
  force = grad(lambda R, *args, **kwargs: -energy_fn(R, *args, **kwargs))
  dt_2 = dt / 2.0
  dt_4 = dt_2 / 2.0
  dt_8 = dt_4 / 2.0
  dt, dt_2, dt_4, dt_8, tau = static_cast(dt, dt_2, dt_4, dt_8, tau)

  def init_fun(key, R):
    V = random.normal(key, R.shape, dtype=R.dtype)
    KE = 0.5 * np.mean(V ** 2)

    # Nose-Hoover parameters.
    xi = np.zeros(chain_length, R.dtype)
    v_xi = np.zeros(chain_length, R.dtype)

    DOF = R.shape[0] * R.shape[1]
    Q = tau ** f32(2) * np.ones(chain_length, dtype=R.dtype)
    Q = ops.index_update(Q, 0, Q[0] * DOF)

    return R, V, KE, xi, v_xi, Q
  def step_chain(KE, V, xi, v_xi, Q, DOF, T):
    """Applies a single update to the chain parameters and rescales velocity."""
    M = chain_length - 1

    G = (Q[M - 1] * v_xi[M - 1] ** f32(2) - T) / Q[M]
    v_xi = ops.index_add(v_xi, M, dt_4 * G)

    G = (f32(2.0) * KE - DOF * T) / Q[0]
    scale = np.exp(-dt_8 * v_xi[1])
    v_xi = ops.index_update(v_xi, 0, scale * (scale * v_xi[0] + dt_4 * G))

    return KE, V, xi, v_xi
  def apply_fun(state, t=0.0, **kwargs):
    R, V, KE, xi, v_xi, Q = state

    DOF = R.shape[0] * R.shape[1]

    KE, V, xi, v_xi = step_chain(KE, V, xi, v_xi, Q, DOF, T)
    R = R + dt_2 * V

    F = force(R, t=t, **kwargs)

    V = V + dt * F
    KE = 0.5 * np.mean(V ** 2)
    R = R + dt_2 * V

    KE, V, xi, v_xi = step_chain(KE, V, xi, v_xi, Q, DOF, T)

    return R, V, KE, xi, v_xi, Q

  return init_fun, apply_fun
```

The following code proceeds without error,
```python
    def do_sim(scale):
      key = random.PRNGKey(0)
      R_key, R0_key, V_key = random.split(key, 3)
      R = random.normal(
        R_key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
      R0 = random.normal(
        R0_key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)

      E = functools.partial(
        lambda R, R0, **kwargs: scale * np.sum((R - R0) ** 2), R0=R0)

      init_fn, apply_fn = nvt_nose_hoover(E, 1e-3, 0.1)

      state = init_fn(V_key, R)

      for _ in range(2):
        state = apply_fn(state)

      return E(state[0])
    assert grad(do_sim)(1.0) > 0.0
```

However, when the apply_fn is jitted, as in
```python
    def do_sim(scale):
      key = random.PRNGKey(0)
      R_key, R0_key, V_key = random.split(key, 3)
      R = random.normal(
        R_key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
      R0 = random.normal(
        R0_key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)

      _, shift = space.free()

      E = functools.partial(
        lambda R, R0, **kwargs: scale * np.sum((R - R0) ** 2), R0=R0)

      init_fn, apply_fn = simulate.nvt_nose_hoover(E, 1e-3, 0.1)
      apply_fn = jit(apply_fn)

      state = init_fn(V_key, R)

      for _ in range(2):
        state = apply_fn(state)

      return E(state[0])
    assert grad(do_sim)(1.0) > 0.0
```
The code throws an error with the following stack trace,
```
Traceback (most recent call last):
  File "/usr/local/google/home/schsam/.local/lib/python2.7/site-packages/absl/third_party/unittest3_backport/case.py", line 37, in testPartExecutor
    yield
  File "/usr/local/google/home/schsam/.local/lib/python2.7/site-packages/absl/third_party/unittest3_backport/case.py", line 162, in run
    testMethod()
  File "/usr/local/google/home/schsam/.local/lib/python2.7/site-packages/absl/testing/parameterized.py", line 262, in bound_param_test
    test_method(self, **testcase_params)
  File "simulate_test.py", line 164, in test_grad_through_nvt
    assert grad(do_sim)(1.0) > 0.0
  File "/usr/local/google/home/schsam/Source/jax/jax/api.py", line 235, in grad_f
    _, g = value_and_grad_f(*args, **kwargs)
  File "/usr/local/google/home/schsam/Source/jax/jax/api.py", line 289, in value_and_grad_f
    g = vjp_py(onp.ones((), dtype=dtype))
  File "/usr/local/google/home/schsam/Source/jax/jax/api_util.py", line 62, in apply_jaxtree_fun
    ans = fun(*args)
  File "/usr/local/google/home/schsam/Source/jax/jax/api.py", line 822, in out_vjp_packed
    return out_vjp(cotangent_in)
  File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/ad.py", line 112, in vjp_
    _, arg_cts = backward_pass(jaxpr, consts, (), dummy_args, dummy_primal_and_ct)
  File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/ad.py", line 180, in backward_pass
    eqn.params, subjaxprs, sub_consts, sub_freevar_vals, invals, ct_in)
  File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/ad.py", line 536, in call_transpose
    ans = primitive.bind(fun, all_args, **params)
  File "/usr/local/google/home/schsam/Source/jax/jax/core.py", line 636, in call_bind
    ans = primitive.impl(f, *args, **params)
  File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/xla.py", line 591, in xla_call_impl
    compiled_fun = xla_callable(fun, device_values, *map(abstractify, args))
  File "/usr/local/google/home/schsam/Source/jax/jax/linear_util.py", line 208, in memoized_fun
    ans = call(f, *args)
  File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/xla.py", line 604, in xla_callable
    jaxpr, (pval, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
  File "/usr/local/google/home/schsam/Source/jax/jax/linear_util.py", line 147, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/ad.py", line 186, in backward_pass
    cts_out = get_primitive_transpose(eqn.primitive)(ct_in, *invals, **eqn.params)
  File "/usr/local/google/home/schsam/Source/jax/jax/lax/lax.py", line 2818, in _scatter_add_transpose_rule
    for i in xrange(len(t.shape)):
AttributeError: 'Zero' object has no attribute 'shape'
```


Discussion:
Looks like just a bug in the scatter transpose rule; it doesn’t handle symbolic zeros like transpose rules need to. 
Totally for my own curiosity, is there a reason why this only affects the version where apply_fn was jitted?

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -2804,6 +2804,9 @@ def _scatter_add_transpose_rule(t, operand, scatter_indices, updates,
                                 update_jaxpr, update_consts, dimension_numbers,
                                 updates_shape):
   assert scatter_indices is not None
+  if t is ad_util.zero:
+    return [ad_util.zero, None, ad_util.zero]
+
   operand_t = update_t = None
   if operand is None:
     operand_t = t


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -30,6 +30,7 @@
 
 import numpy as onp
 
+import jax.ops
 from jax import api
 from jax import lax
 from jax import numpy as lnp
@@ -1552,6 +1553,17 @@ def testIssue764(self):
                           7.66067839e-174], onp.float64)
     self.assertAllClose(f(x), expected, check_dtypes=False)
 
+  def testIssue776(self):
+    """Tests that the scatter-add transpose rule instantiates symbolic zeros."""
+    def f(u):
+      y = jax.ops.index_add(onp.ones(10,), [2, 4, 5], u)
+      # The transpose rule for lax.tie_in returns a symbolic zero for its first
+      # argument.
+      return lax.tie_in(y, 7.)
+
+    self.assertAllClose(onp.zeros(3,), api.grad(f)(onp.ones(3,)),
+                        check_dtypes=True)
+
 
 if __name__ == "__main__":
   absltest.main()