Repository: google/jax

Issue:
segment_sum primitives / advanced indexing in jax.ops.index_add
Would it be useful to others (aside from me) to have better support of sorted and unsorted segment_sums?

https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_sum
https://www.tensorflow.org/api_docs/python/tf/math/segment_sum

In numpy one way to do unsorted segment sums is to sort then call np.add.reduceat, but this doesn't seem to be in jax or autograd:

```
>>> import jax.numpy as np
np.add
>>> np.add
<function _one_to_one_binop.<locals>.<lambda> at 0x1430cdbf8>
>>> np.add.reduceat
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'function' object has no attribute 'reduceat'
```






Discussion:
Does `jax.ops.index_add` do what you need? It seems roughly equivalent to `unsorted_segment_sum`.

https://jax.readthedocs.io/en/latest/_autosummary/jax.ops.index_add.html#jax.ops.index_add
Ah, I guess it does not *yet* because it does not support advanced indexing. We should fix that.
Nice! Would it be hard to support the basic operators +,-,*,/ on top of add? Not sure if there's associativity assumptions here in the underlying implementation or not

Edit: Nevermind: pretty sure - and / will be very hard to do correctly now that I think about it a little more for more than a single operation (since JAX says that it applies all the updates)
After playing around a little bit with the numpy version of advanced index assignment, I take it the idea is to implement the unsorted/sorted segment sums using something similar to:
``` python
ys = np.arange(5)
idxs = [0,0,1,0,1]
sums = np.zeros(2) # two slots
seg_sum = jax.ops.index_add(sums, jax.ops.index[idxs], ys) 
```

edit: of course the jax version will work because it actually accumulates properly as opposed to just applying the last element


Yup, that's the idea!

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -2153,9 +2153,7 @@ def _is_slice_none(idx):
 def _is_advanced_int_indexer(idx):
   """Returns True if idx should trigger int array indexing, False otherwise."""
   # https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#advanced-indexing
-  if isinstance(idx, (tuple, list)):
-    # We assume this check comes *after* the check for non-advanced tuple index,
-    # and hence we already know at least one element is a sequence if it's a tuple
+  if isinstance(idx, (tuple, list)) and _any(onp.ndim(elt) != 0 for elt in idx):
     return _all(e is None or e is Ellipsis or isinstance(e, slice)
                 or _is_int_arraylike(e) for e in idx)
   else:
diff --git a/jax/ops/__init__.py b/jax/ops/__init__.py
--- a/jax/ops/__init__.py
+++ b/jax/ops/__init__.py
@@ -14,4 +14,4 @@
 
 from __future__ import absolute_import
 
-from .scatter import index, index_add, index_update
\ No newline at end of file
+from .scatter import index, index_add, index_update, segment_sum
diff --git a/jax/ops/scatter.py b/jax/ops/scatter.py
--- a/jax/ops/scatter.py
+++ b/jax/ops/scatter.py
@@ -26,6 +26,19 @@
 from ..numpy import lax_numpy as np
 
 
+# TODO(mattjj): clean up this logic
+def _is_advanced_int_indexer(idx):
+  _int = lambda aval: not aval.shape and onp.issubdtype(aval.dtype, onp.integer)
+  try:
+    abstract_idx = core.get_aval(idx)
+  except TypeError:
+    abstract_idx = None
+  out = not (isinstance(abstract_idx, ConcreteArray) and _int(abstract_idx) or
+             isinstance(abstract_idx, ShapedArray) and _int(abstract_idx) or
+             isinstance(idx, slice) or
+             isinstance(idx, tuple) and all(onp.ndim(elt) == 0 for elt in idx))
+  return out and np._is_advanced_int_indexer(idx)
+
 def _scatter_update(x, idx, y, scatter_op):
   """Helper for indexed updates.
 
@@ -33,11 +46,19 @@ def _scatter_update(x, idx, y, scatter_op):
     x[idx] op= y
   except in a pure functional way, with no in-place updating.
 
-  Support NumPy-style basic indexing only, i.e., `idx` must be
-  `None`, an integer, a `slice` object, or ellipses, or a tuple of the above.
+  Args:
+    x: ndarray to be updated.
+    idx: None, an integer, a slice, an ellipsis, an ndarray with integer dtype,
+      or a tuple of those indicating the locations of `x` into which to scatter-
+      update the values in `y`.
+    y: values to be scattered.
+    scatter_op: callable, either lax.scatter or lax.scatter_add.
 
-  TODO(phawkins): support advanced indexing.
+  Returns:
+    An ndarray representing an updated `x` after performing the scatter-update.
   """
+  # For more clues on the logic of this implementation, see the code for
+  # jax.numpy._rewriting_take (which has links to NumPy docs).
 
   x = np.asarray(x)
   y = np.asarray(y)
@@ -45,14 +66,45 @@ def _scatter_update(x, idx, y, scatter_op):
   y_shape = np.shape(y)
   y = lax.convert_element_type(y, lax.dtype(x))
 
+  # Check if there's advanced indexing going on, and handle differently based on
+  # whether it is or isn't mixed with basic indexing.
+  if _is_advanced_int_indexer(idx):
+    if np._is_advanced_int_indexer_without_slices(idx):
+      if isinstance(idx, (tuple, list)):
+        if any(onp.shape(e) for e in idx):
+          # At least one sequence element in the index list means broadcasting.
+          idx = np.broadcast_arrays(*idx)
+        else:
+          # The index list is a flat list of integers.
+          idx = [lax.concatenate([lax.reshape(e, (1,)) for e in idx], 0)]
+      else:
+        # The indexer is just a single integer array.
+        idx = [idx]
+
+      stacked_idx = np.concatenate(
+          [np.mod(np.reshape(a, (-1, 1)), np._constant_like(a, x.shape[i]))
+          for i, a in enumerate(idx)], axis=1)
+
+      y = np.broadcast_to(y, idx[0].shape + onp.shape(x)[len(idx):])
+      y = lax.reshape(y, (stacked_idx.shape[0],) + onp.shape(x)[len(idx):])
+
+      dnums = lax.ScatterDimensionNumbers(
+          update_window_dims=tuple(range(1, y.ndim)),
+          inserted_window_dims=tuple(range(len(idx))),
+          scatter_dims_to_operand_dims=tuple(range(len(idx))))
+      return scatter_op(x, stacked_idx, y, dnums)
+    elif np._is_advanced_int_indexer(idx):
+      # TODO(mattjj, phawkins): one of us is going to implement this case someday
+      msg = "Unimplemented case for indexed update. Open a feature request!"
+      raise NotImplementedError(msg)
+    else:
+      assert False  # unreachable
+
+  # At this point there's no advanced indexing going on, so we process each
+  # element of the index one at a time to build up a scatter.
   if not isinstance(idx, tuple):
     idx = (idx,)
 
-  # Test for unsupported advanced indexing and report an error.
-  if any(onp.ndim(elt) != 0 for elt in idx):
-    raise NotImplementedError("Unimplemented case for indexed update. Advanced "
-                              "indexing is not yet implemented.")
-
   # Remove ellipses and add trailing slice(None)s.
   idx = np._canonicalize_tuple_index(x, idx)
 
@@ -189,10 +241,11 @@ def index_add(x, idx, y):
   (e.g., due to concurrency on some hardware platforms).
 
   Args:
-    x: an array.
-    idx: a Numpy-style basic index, consisting of `None`, integers, `slice`
-      objects, ellipses, or a tuple of the above. A convenient syntactic sugar
-      for forming indices is via the :data:`jax.ops.index` object.
+    x: an array with the values to be updated.
+    idx: a Numpy-style index, consisting of `None`, integers, `slice` objects,
+      ellipses, ndarrays with integer dtypes, or a tuple of the above. A
+      convenient syntactic sugar for forming indices is via the
+      :data:`jax.ops.index` object.
     y: the array of updates. `y` must be broadcastable to the shape of the
       array that would be returned by `x[idx]`.
 
@@ -225,10 +278,11 @@ def index_update(x, idx, y):
   updates on some hardware platforms).
 
   Args:
-    x: an array.
-    idx: a Numpy-style basic index, consisting of `None`, integers, `slice`
-      objects, ellipses, or a tuple of the above. A convenient syntactic sugar
-      for forming indices is via the :data:`jax.ops.index` object.
+    x: an array with the values to be updated.
+    idx: a Numpy-style index, consisting of `None`, integers, `slice` objects,
+      ellipses, ndarrays with integer dtypes, or a tuple of the above. A
+      convenient syntactic sugar for forming indices is via the
+      :data:`jax.ops.index` object.
     y: the array of updates. `y` must be broadcastable to the shape of the
       array that would be returned by `x[idx]`.
 
@@ -244,3 +298,32 @@ def index_update(x, idx, y):
          [1., 1., 1., 6., 6., 6.]], dtype=float32)
   """
   return _scatter_update(x, idx, y, lax.scatter)
+
+def segment_sum(data, segment_ids, num_segments=None):
+  """Computes the sum within segments of an array.
+
+  Similar to TensorFlow's segment_sum:
+  https://www.tensorflow.org/api_docs/python/tf/math/segment_sum
+
+  Args:
+    data: an array with the values to be summed.
+    segment_ids: an array with integer dtype that indicates the segments of
+      `data` (along its leading axis) to be summed. Values can be repeated and
+      need not be sorted. Values outside of the range [0, num_segments) are
+      wrapped into that range by applying np.mod.
+    num_segments: optional, an int with positive value indicating the number of
+      segments. The default is ``max(segment_ids % data.shape[0]) + 1`` but
+      since `num_segments` determines the size of the output, a static value
+      must be provided to use `segment_sum` in a `jit`-compiled function.
+
+  Returns:
+    An array with shape ``(num_segments,) + data.shape[1:]`` representing the
+    segment sums.
+  """
+  if num_segments is None:
+    num_segments = np.max(np.mod(segment_ids, data.shape[0])) + 1
+  num_segments = int(num_segments)
+
+  out = np.zeros((num_segments,) + data.shape[1:], dtype=data.dtype)
+  segment_ids = np.mod(segment_ids, num_segments)
+  return index_add(out, segment_ids, data)


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -216,6 +216,157 @@ def check_grads(f, args, order, atol=None, rtol=None, eps=None):
     #   ]),
 ]
 
+ADVANCED_INDEXING_TESTS = [
+    ("One1DIntArrayIndex",
+     [IndexSpec(shape=(3,), indexer=onp.array([0, 1])),
+     IndexSpec(shape=(3, 3), indexer=onp.array([1, 2, 1])),
+     IndexSpec(shape=(3, 4, 5), indexer=onp.array([0, 2, 0, 1])),
+     IndexSpec(shape=(3,), indexer=onp.array([-1, 1])),
+     IndexSpec(shape=(3,), indexer=onp.array([-2, -1])),
+     ]),
+    ("One2DIntArrayIndex",
+     [IndexSpec(shape=(3,), indexer=onp.array([[0, 0]])),
+     IndexSpec(shape=(3, 3), indexer=onp.array([[1, 2, 1],
+                                                [0, 1, -1]])),
+     IndexSpec(shape=(3, 4, 5), indexer=onp.array([[0, 2, 0, 1],
+                                                   [-1, -2, 1, 0]])),
+     ]),
+    ("Two1DIntArrayIndicesNoBroadcasting",
+     [IndexSpec(shape=(3, 3), indexer=[onp.array([0, 1]),
+                                       onp.array([1, 2])]),
+     IndexSpec(shape=(3, 4, 5), indexer=[onp.array([0, 2, 0, 1]),
+                                         onp.array([-1, 0, -1, 2])]),
+     ]),
+    ("Two1DIntArrayIndicesWithBroadcasting",
+     [IndexSpec(shape=(3, 3), indexer=[onp.array([[0, 1]]),
+                                       onp.array([1, 2])]),
+     IndexSpec(shape=(3, 4, 5), indexer=[onp.array([[0, 2, 0, 1]]),
+                                         onp.array([-1, 0, -1, 2])]),
+     ]),
+    ("ListOfPythonInts",
+     [IndexSpec(shape=(3,), indexer=[0, 1, 0]),
+     IndexSpec(shape=(3, 4, 5), indexer=[0, -1]),
+     ]),
+    ("ListOfListsOfPythonInts",
+     [IndexSpec(shape=(3, 4, 5), indexer=[[0, 1]]),
+     IndexSpec(shape=(3, 4, 5), indexer=[[[0], [-1]], [[2, 3, 0, 3]]]),
+     ]),
+    ("TupleOfListsOfPythonInts",
+     [IndexSpec(shape=(3, 4, 5), indexer=([0, 1])),
+     IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0, 3]])),
+     ]),
+    ("ListOfPythonIntsAndIntArrays",
+     [IndexSpec(shape=(3, 4, 5), indexer=[0, onp.array([0, 1])]),
+     IndexSpec(shape=(3, 4, 5), indexer=[0, 1,
+                                         onp.array([[2, 3, 0, 3]])]),
+     ]),
+    ("ListOfListsOfPythonIntsAndIntArrays",
+     [IndexSpec(shape=(3, 4, 5), indexer=[[0, 1], onp.array([0])]),
+     IndexSpec(shape=(3, 4, 5), indexer=[[[0], [-1]],
+                                         onp.array([[2, 3, 0, 3]])]),
+     ]),
+]
+
+ADVANCED_INDEXING_TESTS_NO_REPEATS = [
+    ("One1DIntArrayIndex",
+     [IndexSpec(shape=(3,), indexer=onp.array([0, 1])),
+      IndexSpec(shape=(3, 3), indexer=onp.array([1, 2, 0])),
+      IndexSpec(shape=(3, 4, 5), indexer=onp.array([0, 2, 1])),
+      IndexSpec(shape=(3,), indexer=onp.array([-1, 1])),
+      IndexSpec(shape=(3,), indexer=onp.array([-2, -1])),
+     ]),
+    ("One2DIntArrayIndex",
+     [IndexSpec(shape=(3,), indexer=onp.array([[0, 1]])),
+      IndexSpec(shape=(6, 6), indexer=onp.array([[1, 2, 0],
+                                                 [3, 4, -1]])),
+     ]),
+    ("Two1DIntArrayIndicesNoBroadcasting",
+     [IndexSpec(shape=(3, 3), indexer=[onp.array([0, 1]),
+                                       onp.array([1, 2])]),
+      IndexSpec(shape=(4, 5, 6), indexer=[onp.array([0, 2, 1, 3]),
+                                          onp.array([-1, 0, -2, 1])]),
+     ]),
+    ("Two1DIntArrayIndicesWithBroadcasting",
+     [IndexSpec(shape=(3, 3), indexer=[onp.array([[0, 1]]),
+                                       onp.array([1, 2])]),
+      IndexSpec(shape=(4, 5, 6), indexer=[onp.array([[0, 2, -1, 1]]),
+                                          onp.array([-1, 0, -2, 2])]),
+     ]),
+    ("ListOfPythonInts",
+     [IndexSpec(shape=(3,), indexer=[0, 2, 1]),
+      IndexSpec(shape=(3, 4, 5), indexer=[0, -1]),
+     ]),
+    ("ListOfListsOfPythonInts",
+     [IndexSpec(shape=(3, 4, 5), indexer=[[0, 1]]),
+      IndexSpec(shape=(3, 4, 5), indexer=[[[0], [-1]], [[2, 3, 0]]]),
+     ]),
+    ("TupleOfListsOfPythonInts",
+     [IndexSpec(shape=(3, 4, 5), indexer=([0, 1])),
+      IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0]])),
+     ]),
+    ("ListOfPythonIntsAndIntArrays",
+     [IndexSpec(shape=(3, 4, 5), indexer=[0, onp.array([0, 1])]),
+      IndexSpec(shape=(3, 4, 5), indexer=[0, 1,
+                                          onp.array([[2, 3, 0]])]),
+     ]),
+    ("ListOfListsOfPythonIntsAndIntArrays",
+     [IndexSpec(shape=(3, 4, 5), indexer=[[0, 1], onp.array([0])]),
+      IndexSpec(shape=(3, 4, 5), indexer=[[[0], [-1]],
+                                          onp.array([[2, 3, 0]])]),
+     ]),
+]
+
+MIXED_ADVANCED_INDEXING_TESTS = [
+    ("SlicesAndOneIntArrayIndex",
+     [IndexSpec(shape=(2, 3), indexer=(onp.array([0, 1]), slice(1, 2))),
+     IndexSpec(shape=(2, 3), indexer=(slice(0, 2),
+                                      onp.array([0, 2]))),
+     IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
+                                         onp.array([0, 2]),
+                                         slice(None))),
+     IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
+                                         onp.array([[0, 2], [1, 1]]),
+                                         slice(None))),
+     ]),
+    ("SlicesAndTwoIntArrayIndices",
+     [IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
+                                          onp.array([0, 2]),
+                                          onp.array([-1, 2]))),
+     IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
+                                         Ellipsis,
+                                         onp.array([-1, 2]))),
+     IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
+                                         onp.array([-1, 2]),
+                                         Ellipsis)),
+     IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
+                                         onp.array([-1, 2]),
+                                         slice(1, 3))),
+     IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
+                                         slice(1, 3),
+                                         onp.array([-1, 2]))),
+     IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2, -2]),
+                                         slice(None, None, 2),
+                                         onp.array([-1, 2, -1]))),
+     IndexSpec(shape=(3, 4, 5), indexer=(onp.array([[0, 2], [2, 0]]),
+                                         Ellipsis,
+                                         onp.array([[1, 0], [1, 0]]))),
+     ]),
+    ("NonesAndIntArrayIndices",
+     [IndexSpec(shape=(3, 4, 5), indexer=[onp.array([0, 2]),
+                                          None,
+                                          onp.array([-1, 2])]),
+     IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
+                                         None,
+                                         None,
+                                         onp.array([-1, 2]))),
+     IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
+                                         onp.array([0, 2]),
+                                         None,
+                                         None,
+                                         onp.array([-1, 2]))),
+     ]),
+]
+
 class IndexingTest(jtu.JaxTestCase):
   """Tests for Numpy indexing translation rules."""
 
@@ -371,56 +522,7 @@ def fun(unpacked_indexer, x):
       {"testcase_name": "{}_inshape={}_indexer={}"
        .format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
        "shape": shape, "dtype": dtype, "rng": rng, "indexer": indexer}
-      for name, index_specs in [
-          ("One1DIntArrayIndex",
-           [IndexSpec(shape=(3,), indexer=onp.array([0, 1])),
-            IndexSpec(shape=(3, 3), indexer=onp.array([1, 2, 1])),
-            IndexSpec(shape=(3, 4, 5), indexer=onp.array([0, 2, 0, 1])),
-            IndexSpec(shape=(3,), indexer=onp.array([-1, 1])),
-            IndexSpec(shape=(3,), indexer=onp.array([-2, -1])),
-            ]),
-          ("One2DIntArrayIndex",
-           [IndexSpec(shape=(3,), indexer=onp.array([[0, 0]])),
-            IndexSpec(shape=(3, 3), indexer=onp.array([[1, 2, 1],
-                                                       [0, 1, -1]])),
-            IndexSpec(shape=(3, 4, 5), indexer=onp.array([[0, 2, 0, 1],
-                                                          [-1, -2, 1, 0]])),
-            ]),
-          ("Two1DIntArrayIndicesNoBroadcasting",
-           [IndexSpec(shape=(3, 3), indexer=[onp.array([0, 1]),
-                                             onp.array([1, 2])]),
-            IndexSpec(shape=(3, 4, 5), indexer=[onp.array([0, 2, 0, 1]),
-                                                onp.array([-1, 0, -1, 2])]),
-            ]),
-          ("Two1DIntArrayIndicesWithBroadcasting",
-           [IndexSpec(shape=(3, 3), indexer=[onp.array([[0, 1]]),
-                                             onp.array([1, 2])]),
-            IndexSpec(shape=(3, 4, 5), indexer=[onp.array([[0, 2, 0, 1]]),
-                                                onp.array([-1, 0, -1, 2])]),
-            ]),
-          ("ListOfPythonInts",
-           [IndexSpec(shape=(3,), indexer=[0, 1, 0]),
-            IndexSpec(shape=(3, 4, 5), indexer=[0, -1]),
-            ]),
-          ("ListOfListsOfPythonInts",
-           [IndexSpec(shape=(3, 4, 5), indexer=[[0, 1]]),
-            IndexSpec(shape=(3, 4, 5), indexer=[[[0], [-1]], [[2, 3, 0, 3]]]),
-            ]),
-          ("TupleOfListsOfPythonInts",
-           [IndexSpec(shape=(3, 4, 5), indexer=([0, 1])),
-            IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0, 3]])),
-            ]),
-          ("ListOfPythonIntsAndIntArrays",
-           [IndexSpec(shape=(3, 4, 5), indexer=[0, onp.array([0, 1])]),
-            IndexSpec(shape=(3, 4, 5), indexer=[0, 1,
-                                                onp.array([[2, 3, 0, 3]])]),
-            ]),
-          ("ListOfListsOfPythonIntsAndIntArrays",
-           [IndexSpec(shape=(3, 4, 5), indexer=[[0, 1], onp.array([0])]),
-            IndexSpec(shape=(3, 4, 5), indexer=[[[0], [-1]],
-                                                onp.array([[2, 3, 0, 3]])]),
-            ]),
-      ]
+      for name, index_specs in ADVANCED_INDEXING_TESTS
       for shape, indexer in index_specs
       for dtype in all_dtypes
       for rng in [jtu.rand_default()])
@@ -492,56 +594,7 @@ def testAdvancedIntegerIndexingGrads(self, shape, dtype, rng, indexer):
       {"testcase_name": "{}_inshape={}_indexer={}"
        .format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
        "shape": shape, "dtype": dtype, "rng": rng, "indexer": indexer}
-      for name, index_specs in [
-          ("SlicesAndOneIntArrayIndex",
-           [IndexSpec(shape=(2, 3), indexer=(onp.array([0, 1]), slice(1, 2))),
-            IndexSpec(shape=(2, 3), indexer=(slice(0, 2),
-                                             onp.array([0, 2]))),
-            IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
-                                                onp.array([0, 2]),
-                                                slice(None))),
-            IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
-                                                onp.array([[0, 2], [1, 1]]),
-                                                slice(None))),
-            ]),
-          ("SlicesAndTwoIntArrayIndices",
-           [IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
-                                                onp.array([0, 2]),
-                                                onp.array([-1, 2]))),
-            IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
-                                                Ellipsis,
-                                                onp.array([-1, 2]))),
-            IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
-                                                onp.array([-1, 2]),
-                                                Ellipsis)),
-            IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
-                                                onp.array([-1, 2]),
-                                                slice(1, 3))),
-            IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
-                                                slice(1, 3),
-                                                onp.array([-1, 2]))),
-            IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2, -2]),
-                                                slice(None, None, 2),
-                                                onp.array([-1, 2, -1]))),
-            IndexSpec(shape=(3, 4, 5), indexer=(onp.array([[0, 2], [2, 0]]),
-                                                Ellipsis,
-                                                onp.array([[1, 0], [1, 0]]))),
-            ]),
-          ("NonesAndIntArrayIndices",
-           [IndexSpec(shape=(3, 4, 5), indexer=[onp.array([0, 2]),
-                                                None,
-                                                onp.array([-1, 2])]),
-            IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
-                                                None,
-                                                None,
-                                                onp.array([-1, 2]))),
-            IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
-                                                onp.array([0, 2]),
-                                                None,
-                                                None,
-                                                onp.array([-1, 2]))),
-            ]),
-      ]
+      for name, index_specs in MIXED_ADVANCED_INDEXING_TESTS
       for shape, indexer in index_specs
       for dtype in all_dtypes
       for rng in [jtu.rand_default()])
@@ -706,6 +759,35 @@ def onp_fn(x, y):
     self._CheckAgainstNumpy(onp_fn, jax_fn, args_maker, check_dtypes=True)
     self._CompileAndCheck(jax_fn, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list({
+      "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}".format(
+          name, jtu.format_shape_dtype_string(shape, dtype), indexer,
+          jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
+       "shape": shape, "dtype": dtype, "rng": rng, "indexer": indexer,
+       "update_shape": update_shape, "update_dtype": update_dtype,
+       "op": op
+  } for name, index_specs in ADVANCED_INDEXING_TESTS_NO_REPEATS
+    for shape, indexer in index_specs
+    for op in [UpdateOps.UPDATE, UpdateOps.ADD]
+    for dtype in (all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
+    for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
+    for update_dtype in ([dtype] if op == UpdateOps.ADD else all_dtypes)
+    for rng in [jtu.rand_default()]))
+  def testAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
+                           rng, indexer, op):
+    args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
+    def onp_fn(x, y):
+      x = x.copy()
+      if op == UpdateOps.UPDATE:
+        x[indexer] = y
+      else:
+        x[indexer] += y
+      return x
+
+    jax_op = ops.index_update if op == UpdateOps.UPDATE else ops.index_add
+    jax_fn = lambda x, y: jax_op(x, indexer, y)
+    self._CheckAgainstNumpy(onp_fn, jax_fn, args_maker, check_dtypes=True)
+    self._CompileAndCheck(jax_fn, args_maker, check_dtypes=True)
 
   @parameterized.named_parameters(jtu.cases_from_list({
       "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}".format(
@@ -729,6 +811,32 @@ def testStaticIndexingGrads(self, shape, dtype, update_shape, update_dtype,
     y = rng(update_shape, update_dtype)
     check_grads(jax_fn, (x, y), 2, rtol=1e-3, atol=1e-3, eps=1.)
 
+  def testSegmentSumBehavior(self):
+    # testAdvancedIndexing compares against NumPy, and as a result doesn't check
+    # repeated indices. This test is just a simple manual check, based on
+    # https://www.tensorflow.org/api_docs/python/tf/math/segment_sum
+    data = onp.array([5, 1, 7, 2, 3, 4, 1, 3])
+    segment_ids = onp.array([0, 0, 0, 1, 2, 2, 3, 3])
+
+    ans = ops.index_add(onp.zeros(onp.max(segment_ids) + 1), segment_ids, data)
+    expected = onp.array([13, 2, 7, 4])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testSegmentSum(self):
+    data = onp.array([5, 1, 7, 2, 3, 4, 1, 3])
+    segment_ids = onp.array([0, 0, 0, 1, 2, 2, 3, 3])
+
+    # test with explicit num_segments
+    ans = ops.segment_sum(data, segment_ids, num_segments=4)
+    expected = onp.array([13, 2, 7, 4])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+    # test without explicit num_segments
+    ans = ops.segment_sum(data, segment_ids)
+    expected = onp.array([13, 2, 7, 4])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+
 
 if __name__ == "__main__":
   absltest.main()