fix higher-order rev-mode lax.index_take test failures 
With the new-and-more-thorough `check_grads`, we started getting some failures in higher-order tests of `lax.index_take`. See commented-out line in lax_test.py's testIndexTake.
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 56b96e0b2..5feaa7924 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -662,6 +662,7 @@ def scatter(operand, scatter_indices, updates, dimension_numbers):
 
 def index_take(src, idxs, axes):
   indices = concatenate([reshape(i, [i.shape[0], 1]) for i in idxs], 1)
+  indices = indices % onp.array([src.shape[ax] for ax in axes])
   slice_sizes = list(src.shape)
   for ax in axes:
     slice_sizes[ax] = 1
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 4a5f47cef..184882646 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2038,12 +2038,9 @@ class LaxAutodiffTest(jtu.JaxTestCase):
       ]
       for rng in [jtu.rand_default()]))
   def testIndexTakeGrad(self, shape, dtype, idxs, axes, rng):
-    idxs = tuple(rng(e.shape, e.dtype) for e in idxs)
     src = rng(shape, dtype)
     index_take = lambda src: lax.index_take(src, idxs, axes)
-    check_grads(index_take, (src,), 2, ["fwd"], 1e-2, 1e-2, 1)
-    # TODO(mattjj): fix rev mode failures here!
-    # check_grads(index_take, (src,), 2, ["fwd", "rev"], 1e-2, 1e-2, 1)
+    check_grads(index_take, (src,), 2, ["fwd", "rev"], 1e-2, 1e-2, 1)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_idxs={}_dnums={}_slice_sizes={}".format(
['tests/lax_test.py', 'jax/lax/lax.py']