Repository: google/jax

Issue:
NotImplementedError: Batching rule for 'select_and_gather_add' not implemented
I'm trying to compute per-example hessian-vector products, where the model includes a MaxPool layer, and I'm getting this error:
```
  File "cnn.py", line 62, in loss
    logits = predict_fun(params, inputs)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/experimental/stax.py", line 295, in apply_fun
    inputs = fun(param, inputs, rng=rng, **kwargs)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/experimental/stax.py", line 186, in apply_fun
    out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/lax/lax.py", line 949, in reduce_window
    return monoid_reducer(operand, window_dimensions, window_strides, padding)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/lax/lax.py", line 984, in _reduce_window_max
    window_strides=tuple(window_strides), padding=padding)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/core.py", line 155, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/interpreters/ad.py", line 222, in process_primitive
    primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/interpreters/ad.py", line 320, in standard_jvp
    val_out = primitive.bind(*primals, **params)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/core.py", line 155, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/interpreters/ad.py", line 222, in process_primitive
    primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/interpreters/ad.py", line 321, in standard_jvp
    tangents_out = [rule(t, *primals, **params) for rule, t in zip(jvprules, tangents)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/interpreters/ad.py", line 322, in <listcomp>
    if rule is not None and t is not zero]
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/lax/lax.py", line 3590, in _reduce_window_chooser_jvp_rule
    window_strides, padding)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/lax/lax.py", line 1013, in _select_and_gather_add
    window_strides=tuple(window_strides), padding=padding)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/core.py", line 155, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/interpreters/batching.py", line 116, in process_primitive
    batched_primitive = get_primitive_batcher(primitive)
  File "/h/choidami/software/anaconda/lib/python3.7/site-packages/jax/interpreters/batching.py", line 168, in get_primitive_batcher
    raise NotImplementedError(msg.format(p))
NotImplementedError: Batching rule for 'select_and_gather_add' not implemented
```

The code I used is a simple modification of the resnet50.py example in the repo. 
```
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from functools import partial

import numpy.random as npr

from six.moves import xrange

import jax.numpy as np
from jax.config import config
from jax import jit, grad, random, vmap, jvp
from jax.experimental import optimizers
from jax.experimental import stax
from jax.experimental.stax import (Conv, Dense, Flatten, GeneralConv,
                                   MaxPool, Relu, LogSoftmax)
from jax.flatten_util import ravel_pytree


def CNN(num_classes):
  return stax.serial(
      stax.Conv(32, (5, 5), strides=(2, 2)),
      Relu,
      stax.MaxPool((2, 2), strides=(2, 2)),
      Flatten,
      Dense(num_classes), LogSoftmax)


if __name__ == "__main__":
    rng_key = random.PRNGKey(0)

    batch_size = 8
    num_classes = 10
    input_shape = (batch_size, 32, 32, 3)
    step_size = 0.1
    num_steps = 10

    init_fun, predict_fun = CNN(num_classes)
    _, init_params = init_fun(rng_key, input_shape)

    def loss(params, batch):
        inputs, targets = batch
        logits = predict_fun(params, inputs)
        return -np.sum(logits * targets)

    def accuracy(params, batch):
        inputs, targets = batch
        target_class = np.argmax(targets, axis=-1)
        predicted_class = np.argmax(predict_fun(params, inputs), axis=-1)
        return np.mean(predicted_class == target_class)

    def synth_batches():
        rng = npr.RandomState(0)
        while True:
            images = rng.rand(*input_shape).astype('float32')
            labels = rng.randint(num_classes, size=(batch_size, 1))
            onehot_labels = labels == np.arange(num_classes)
            yield images, onehot_labels

    opt_init, opt_update, get_params = optimizers.momentum(step_size, mass=0.9)
    batches = synth_batches()


    def fgh_fn(loss, params, batch, v):
        _loss = partial(loss, batch=batch)
        grads, hvp = jvp(grad(_loss), [params], [v])
        return grads, hvp


    def per_example_gh(loss, params, batch, v):
        fgh = partial(fgh_fn, loss)

        # Add extra batch dimension in case some functions assume it.
        def batch_fgh(params, batch, v):
            X, y = batch
            X = np.expand_dims(X, 0)
            y = np.expand_dims(y, 0)
            return fgh(params, (X, y), v)

        return vmap(batch_fgh, in_axes=(None, 0, None))(params, batch, v)

    @jit
    def update(i, opt_state, batch):
        params = get_params(opt_state)
        flat_params, unravel = ravel_pytree(params)

        def flat_loss(flat_params, batch):
            return loss(unravel(flat_params), batch)
        gs, hvps = per_example_gh(flat_loss, flat_params, batch, np.ones_like(flat_params))
        return opt_update(i, grad(loss)(params, batch), opt_state)

    opt_state = opt_init(init_params)
    for i in xrange(num_steps):
        opt_state = update(i, opt_state, next(batches))
    trained_params = get_params(opt_state)
```



Discussion:

Resolving patch:
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -262,6 +262,8 @@ def broadcast(x, sz, axis):
 def moveaxis(x, src, dst):
   if core.get_aval(x) is core.abstract_unit:
     return core.unit
+  if src == dst:
+    return x
   src, dst = src % x.ndim, dst % x.ndim
   perm = [i for i in range(onp.ndim(x)) if i != src]
   perm.insert(dst, src)
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
--- a/jax/lax/__init__.py
+++ b/jax/lax/__init__.py
@@ -16,7 +16,8 @@
 from .lax import *
 from .lax import (_reduce_sum, _reduce_max, _reduce_min, _reduce_or,
                   _reduce_and, _reduce_window_sum, _reduce_window_max,
-                  _reduce_window_min, _reduce_window_prod, _float, _complex,
+                  _reduce_window_min, _reduce_window_prod,
+                  _select_and_gather_add, _float, _complex,
                   _input_dtype, _const, _eq_meet, _safe_mul,
                   _broadcasting_select, _check_user_dtype_supported,
                   _one, _const, _upcast_fp16_for_computation)
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -3778,7 +3778,7 @@ def _select_and_gather_add_translation(
                                          canonicalize_types=False)
 
   if double_word_reduction:
-    # XLA doesn't yet implement ReduceWindow on tuples (Google bug b/73062247), so
+    # TODO(b/73062247): XLA doesn't yet implement ReduceWindow on tuples, so
     # we implement a pair-wise ReduceWindow by packing two k-bit values into
     # 2k-bit unsigned integer using bit tricks.
     word_dtype = _UINT_DTYPES[nbits]
@@ -3852,7 +3852,7 @@ def reducer():
     return c.Build()
 
 
-  assert select_prim is ge_p or select_prim is le_p
+  assert select_prim is ge_p or select_prim is le_p, select_prim
   init = -onp.inf if select_prim is ge_p else onp.inf
   out = c.ReduceWindow(pack(operand, tangents),
                        pack(const(c, dtype, init), const(c, dtype, 0)),
@@ -3885,12 +3885,30 @@ def _select_and_gather_add_transpose(
                                    window_strides, padding)
   return [result, None]
 
+def _select_and_gather_add_batching_rule(
+    batched_args, batch_dims, select_prim, window_dimensions, window_strides,
+    padding):
+  t, x = batched_args
+  t_bdim, x_bdim = batch_dims
+  size = next(a.shape[bdim] for a, bdim in zip(batched_args, batch_dims)
+              if bdim is not None)
+  t = batching.bdim_at_front(t, t_bdim, size)
+  x = batching.bdim_at_front(x, x_bdim, size)
+  window_dimensions = (1,) + window_dimensions
+  window_strides = (1,) + window_strides
+  out = _select_and_gather_add(t, x, select_prim, window_dimensions,
+                               window_strides, padding)
+  return (out, 0)
+
+
 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
+  _select_and_gather_add_transpose
+batching.primitive_batchers[select_and_gather_add_p] = \
+  _select_and_gather_add_batching_rule
 xla.backend_specific_translations['tpu'][select_and_gather_add_p] = partial(
   _select_and_gather_add_translation,
   max_bits=32)


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2894,6 +2894,32 @@ def fun(operand):
       for bdims in all_bdims(shape):
         self._CheckBatching(fun, 3, bdims, (shape,), dtype, rng)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_dtype={}_padding={}".format(onp.dtype(dtype).name,
+                                                      padding),
+       "dtype": dtype, "padding": padding, "rng_factory": rng_factory}
+      for dtype in float_dtypes
+      for padding in ["VALID", "SAME"]
+      for rng_factory in [jtu.rand_small]))
+  def testSelectAndGatherAdd(self, dtype, padding, rng_factory):
+    rng = rng_factory()
+    all_configs = itertools.chain(
+        itertools.product(
+            [(4, 6)],
+            [(2, 1), (1, 2)],
+            [(1, 1), (2, 1), (1, 2)]),
+        itertools.product(
+            [(3, 2, 4, 6)], [(1, 1, 2, 1), (2, 1, 2, 1)],
+            [(1, 2, 2, 1), (1, 1, 1, 1)]))
+
+    def fun(operand, tangents):
+      return lax._select_and_gather_add(operand, tangents, lax.ge_p, dims,
+                                        strides, padding)
+
+    for shape, dims, strides in all_configs:
+      for bdims in all_bdims(shape, shape):
+        self._CheckBatching(fun, 3, bdims, (shape, shape), dtype, rng)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_bdims={}_fft_ndims={}"
        .format(shape, bdims, fft_ndims),