Repository: google/jax

Issue:
Error with jax.grad when using jax.ops.index_update.
(Appears to be different than issue #1104)

Taking the gradient through jax.ops.index_update fails in certain cases. Code to reproduce:

```
def fun():
  def _fun(x):
    n = x.shape[0]
    y = np.array(list(range(n)))
    x = jax.ops.index_update(x, np.diag_indices(n), y)
    return np.sum(x)
  
  x = np.ones((5,5))
  print(_fun(x))   # Succeeds
  return jax.grad(_fun)(x)
fun()
```
Interestingly, the above function succeeds when using `jax.ops.index_update(x, np.diag_indices(n), y - 0 * np.diag(x))`, so my guess is the lack of dependency on `x` is somehow the problem.

Stack trace 
```
<ipython-input-2-176922dbb306> in fun()
      8   x = np.ones((5,5))
      9   print(_fun(x))   # Succeeds
---> 10   return jax.grad(_fun)(x)

/usr/local/lib/python3.6/dist-packages/jax/api.py in grad_f(*args, **kwargs)
    345   @wraps(fun, docstr=docstr, argnums=argnums)
    346   def grad_f(*args, **kwargs):
--> 347     _, g = value_and_grad_f(*args, **kwargs)
    348     return g
    349 

/usr/local/lib/python3.6/dist-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    400     f_partial, dyn_args = _argnums_partial(f, argnums, args)
    401     if not has_aux:
--> 402       ans, vjp_py = vjp(f_partial, *dyn_args)
    403     else:
    404       ans, vjp_py, aux = vjp(f_partial, *dyn_args, has_aux=True)

/usr/local/lib/python3.6/dist-packages/jax/api.py in vjp(fun, *primals, **kwargs)
   1255   if not has_aux:
   1256     flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-> 1257     out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
   1258     out_tree = out_tree()
   1259   else:

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux)
    105 def vjp(traceable, primals, has_aux=False):
    106   if not has_aux:
--> 107     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    108   else:
    109     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     94   _, in_tree = tree_flatten(((primals, primals), {}))
     95   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
---> 96   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
     97   pval_primals, pval_tangents = tree_unflatten(out_tree(), out_pvals)
     98   aval_primals, const_primals = unzip2(pval_primals)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, **kwargs)
    341   with new_master(JaxprTrace) as master:
    342     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 343     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    344     assert not env
    345     del master

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    151 
    152     del gen
--> 153     ans = self.f(*args, **dict(self.params, **kwargs))
    154     del args
    155     while stack:

<ipython-input-2-176922dbb306> in _fun(x)
      3     n = x.shape[0]
      4     y = np.array(list(range(n)))
----> 5     x = jax.ops.index_update(x, np.diag_indices(n), y)
      6     return np.sum(x)
      7 

/usr/local/lib/python3.6/dist-packages/jax/ops/scatter.py in index_update(x, idx, y)
    247          [1., 1., 1., 6., 6., 6.]], dtype=float32)
    248   """
--> 249   return _scatter_update(x, idx, y, lax.scatter)
    250 
    251 def segment_sum(data, segment_ids, num_segments=None):

/usr/local/lib/python3.6/dist-packages/jax/ops/scatter.py in _scatter_update(x, idx, y, scatter_op)
     55   # is more or less a transpose of the gather equivalent.
     56   treedef, static_idx, dynamic_idx = np._split_index_for_jit(idx)
---> 57   return _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx)
     58 
     59 

/usr/local/lib/python3.6/dist-packages/jax/ops/scatter.py in _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx)
     81     scatter_dims_to_operand_dims=indexer.dnums.start_index_map
     82   )
---> 83   return scatter_op(x, indexer.gather_indices, y, dnums)
     84 
     85 

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in scatter(operand, scatter_indices, updates, dimension_numbers)
    840       operand, scatter_indices, updates, update_jaxpr=jaxpr,
    841       update_consts=consts, dimension_numbers=dimension_numbers,
--> 842       updates_shape=updates.shape)
    843 
    844 def index_take(src, idxs, axes):

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, *args, **kwargs)
    153 
    154     tracers = map(top_trace.full_raise, args)
--> 155     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    156     if self.multiple_results:
    157       return map(full_lower, out_tracer)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in process_primitive(self, primitive, tracers, params)
    220           "Forward-mode differentiation rule for '{}' not implemented"
    221           .format(primitive))
--> 222     primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
    223     if primitive.multiple_results:
    224       return [JVPTracer(self, x, t) for x, t in zip(primal_out, tangent_out)]

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in _scatter_jvp(primals, tangents, update_jaxpr, update_consts, dimension_numbers, updates_shape)
   3286                             g_operand, _zeros(g_operand))
   3287   masked_g_updates = select(eq(update_ids, gathered_update_ids),
-> 3288                             g_updates, _zeros(g_updates))
   3289 
   3290   # d) perform a scatter-add to compute the tangent output.

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in select(pred, on_true, on_false)
    648   operator.
    649   """
--> 650   return select_p.bind(pred, on_true, on_false)
    651 
    652 def slice(operand, start_indices, limit_indices, strides=None):

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, *args, **kwargs)
    150     top_trace = find_top_trace(args)
    151     if top_trace is None:
--> 152       return self.impl(*args, **kwargs)
    153 
    154     tracers = map(top_trace.full_raise, args)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    138   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
    139   abstract_args = map(abstractify, args)
--> 140   compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
    141   return compiled_fun(*args)
    142 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in xla_primitive_callable(prim, *abstract_args, **params)
    144 def xla_primitive_callable(prim, *abstract_args, **params):
    145   backend = params.get('backend', None)
--> 146   aval_out = prim.abstract_eval(*abstract_args, **params)
    147   if prim.multiple_results:
    148     handlers = tuple(map(aval_to_result_handler, aval_out))

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in standard_abstract_eval(shape_rule, dtype_rule, *args, **kwargs)
   1486     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
   1487   elif least_specialized is ShapedArray:
-> 1488     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
   1489   elif least_specialized is UnshapedArray:
   1490     return UnshapedArray(dtype_rule(*args, **kwargs))

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in _select_shape_rule(pred, on_true, on_false)
   2573     msg = ("select pred must be scalar or have the same shape as on_true and "
   2574            "on_false, got pred shape {} for on_true and on_false of shape {}.")
-> 2575     raise TypeError(msg.format(pred.shape, on_true.shape))
   2576   return on_true.shape
   2577 

TypeError: select pred must be scalar or have the same shape as on_true and on_false, got pred shape (5,) for on_true and on_false of shape (1, 5).
```


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -3232,6 +3232,9 @@ def _scatter_jvp(primals, tangents, update_jaxpr, update_consts,
     tangent_out = ad_util.zero
     return val_out, tangent_out
 
+  g_operand = ad.instantiate_zeros(operand, g_operand)
+  g_updates = ad.instantiate_zeros(updates, g_updates)
+
   # If there are overlapping indices in the scatter, it is unspecified which
   # update "wins". So we use the following perhaps surprising scheme:
   # a) attach a positive ID to each update in updates, forming (value, id) pairs
@@ -3298,8 +3301,6 @@ def _scatter_jvp(primals, tangents, update_jaxpr, update_consts,
                          slice_sizes=slice_sizes)
 
   # c) mask off input JVP elements that do not correspond to a primal output.
-  g_operand = ad.instantiate_zeros(operand, g_operand)
-  g_updates = ad.instantiate_zeros(updates, g_updates)
   masked_g_operand = select(eq(scattered_ids, _zeros(scattered_ids)),
                             g_operand, _zeros(g_operand))
   masked_g_updates = select(eq(update_ids, gathered_update_ids),


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2465,6 +2465,16 @@ def testScatterGrad(self, arg_shape, dtype, idxs, update_shape, dnums,
     y = rng(update_shape, dtype)
     check_grads(scatter, (x, y), 2, ["fwd", "rev"], 1e-2, 1e-2, 1.)
 
+  def testScatterGradSymbolicZeroUpdate(self):
+    # https://github.com/google/jax/issues/1901
+    def f(x):
+      n = x.shape[0]
+      y = onp.arange(n, dtype=x.dtype)
+      return jax.ops.index_update(x, onp.diag_indices(n), y)
+    rng = jtu.rand_default()
+    check_grads(f, (rng((5, 5), onp.float32),), 2, ["fwd", "rev"], 1e-2, 1e-2,
+                1.)
+
   def testStopGradient(self):
     def f(x):
       return lax.sin(x) * lax.cos(lax.stop_gradient(x))