Repository: jax-ml/jax
Pull request: Instantiate symbolic zeros in the scatter-add transpose rule.
Fixes #776.

Problem statement:
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'
```

Hints from issue 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?

Commit message: Instantiate symbolic zeros in the scatter-add transpose rule.

Fixes #776.

Patch to review:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 442300300a61..522ccdb982bc 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -2809,6 +2809,7 @@ def _scatter_add_transpose_rule(t, operand, scatter_indices, updates,
     operand_t = t
 
   if updates is None:
+    t = ad.instantiate_zeros(operand, t)
     gather_dnums = GatherDimensionNumbers(
       offset_dims=dimension_numbers.update_window_dims,
       collapsed_slice_dims=dimension_numbers.inserted_window_dims,
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 49072113115b..0a113a0bc5e8 100644
--- 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()


Code review comments:
diff_hunk: @@ -2809,6 +2809,7 @@ def _scatter_add_transpose_rule(t, operand, scatter_indices, updates,
     operand_t = t
 
   if updates is None:
+    t = ad.instantiate_zeros(operand, t)
path: jax/lax/lax.py
text: @user1:
Hmm is it possible for `operand` to be None here? I think yes:

```python
grad(lambda x: lax.tie_in(jax.ops.index_add(x, [0], x)[0], 7.))(np.ones(1))
```

I think we might need to add `operand_shape` to the params of `scatter_add`, since otherwise we don't have enough information to instantiate this zero.

That'd be in keeping with the way we've solved this kind of issue in the past, at least. There might be other options.

@author:
Come to think of it, why do we need to instantiate zeros here instead of just propagating them? I did that instead.