Repository: google/jax

Issue:
Forward-mode differentiation rule for 'select_and_scatter_add' not implemented
Ran into a not implemented error when taking second order gradients through a max pool.
Repro:
```
import jax.numpy as np
from jax.api import jvp
from jax import grad
import jax.random
from jax.experimental import stax
import numpy as onp

data = np.array(onp.random.normal(size=[1, 4, 4, 1]))

model = stax.serial(stax.Conv(1, (1,1), (1,1), 'SAME'),
                   stax.MaxPool((2,2), strides=(2,2)))

model_init, model_predict = model

key = jax.random.PRNGKey(0)
params = model_init(key, list(data.shape))[1]


def loss(params, data):
  return np.sum(model_predict(params, data))


def hvp(loss, params, batch, v):
  loss_fn = lambda x: loss(x, batch)
  return jvp(grad(loss_fn), [params], [v])[1]

# Raises NotImplementedError: Forward-mode differentiation rule for 'select_and_scatter_add' not implemented
hvp(loss, params, data, params)
```


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
@@ -3252,13 +3252,30 @@ def _select_and_scatter_add_translation(
   return c.SelectAndScatter(operand, select, window_dimensions, window_strides,
                             padding, source, zero, scatter)
 
+def _select_and_scatter_add_jvp(
+    primals, tangents, select_prim, window_dimensions, window_strides,
+    padding):
+  source, operand = primals
+  g_source, g_operand = tangents
+  val_out = _select_and_scatter_add(
+      source, operand, select_prim, window_dimensions, window_strides,
+      padding)
+  del g_operand
+  if g_source is ad_util.zero:
+    tangent_out = ad_util.zero
+  else:
+    tangent_out = _select_and_scatter_add(
+        g_source, operand, select_prim, window_dimensions,
+        window_strides, padding)
+  return val_out, tangent_out
+
 def _select_and_scatter_add_transpose(
     t, source, operand, select_prim, window_dimensions, window_strides,
     padding):
   assert source is None and operand is not None
-  result = _select_and_gather_add(t, operand, select_prim, window_dimensions,
-                                  window_strides, padding)
-  return [result, None]
+  source_t = _select_and_gather_add(t, operand, select_prim, window_dimensions,
+                                    window_strides, padding)
+  return [source_t, None]
 
 def _select_and_scatter_add_batch_rule(batched_args, batch_dims, **kwargs):
   source, operand = batched_args
@@ -3295,6 +3312,7 @@ def _select_and_scatter_add_batch_rule(batched_args, batch_dims, **kwargs):
     _select_and_scatter_add_translation)
 ad.primitive_transposes[select_and_scatter_add_p] = \
     _select_and_scatter_add_transpose
+ad.primitive_jvps[select_and_scatter_add_p] = _select_and_scatter_add_jvp
 batching.primitive_batchers[select_and_scatter_add_p] = \
     _select_and_scatter_add_batch_rule
 
@@ -3375,6 +3393,23 @@ def _select_and_gather_add_translation(
   out = c.ConvertElementType(out, uint_etype)
   return c.BitcastConvertType(out, etype)
 
+def _select_and_gather_add_jvp(
+    primals, tangents, select_prim, window_dimensions, window_strides,
+    padding):
+  source, operand = primals
+  g_source, g_operand = tangents
+  val_out = _select_and_gather_add(
+      source, operand, select_prim, window_dimensions, window_strides,
+      padding)
+  del g_operand
+  if g_source is ad_util.zero:
+    tangent_out = ad_util.zero
+  else:
+    tangent_out = _select_and_gather_add(
+        g_source, operand, select_prim, window_dimensions,
+        window_strides, padding)
+  return val_out, tangent_out
+
 def _select_and_gather_add_transpose(
     t, tangents, operand, select_prim, window_dimensions, window_strides,
     padding):
@@ -3386,6 +3421,7 @@ def _select_and_gather_add_transpose(
 select_and_gather_add_p = standard_primitive(
     _select_and_gather_add_shape_rule, _input_dtype, 'select_and_gather_add',
     _select_and_gather_add_translation)
+ad.primitive_jvps[select_and_gather_add_p] = _select_and_gather_add_jvp
 ad.primitive_transposes[select_and_gather_add_p] = \
     _select_and_gather_add_transpose
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1945,9 +1945,7 @@ def fun(operand):
         self.assertEqual(onp.unique(operand).size, operand.size,
                          msg="test requires operand elements to be unique.")
       jtu.check_vjp(fun, partial(api.vjp, fun), (operand,), 1e-2, 1e-2, 1e-2)
-
-      # TODO(phawkins): enable both gradients after a jaxlib update.
-      # check_grads(fun, (operand,), 1, 1e-2, 1e-2, 1e-2)
+      check_grads(fun, (operand,), 3, 1e-2, 1e-2, 1e-2)
     # pylint: enable=cell-var-from-loop
 
   # TODO(b/205052657): enable more tests when supported