Repository: google/jax

Issue:
Add index_mul, index_div, index_pow
Currently the only scatter operations implemented as primitives are update, add, min and max. `index_update` has limited support for transposing, which means some custom_jvp rules implemented using `index_update` cannot be automatically converted into vjp form. Extending the indexing operations would fix this, in particular in the fix for issue #2380. 


Discussion:
One third of the way there!
Though it's probably not high priority

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -49,6 +49,7 @@
 from jax import lax
 from jax._src.lax.lax import _device_put_raw
 from jax import ops
+from jax._src.ops import scatter
 from jax._src.util import (partial, unzip2, prod as _prod, subvals, safe_zip,
                            canonicalize_axis as _canonicalize_axis, maybe_named_axis)
 from jax.tree_util import tree_leaves, tree_flatten, tree_map
@@ -5842,14 +5843,17 @@ class _IndexUpdateHelper:
 
   The ``at`` property is syntactic sugar for calling the indexed update functions
   defined in :mod:`jax.ops`, and acts as a pure equivalent of in-place
-  modificatons. For further information, see `Syntactic Sugar for Index Update Operators
-  <https://jax.readthedocs.io/en/latest/jax.ops.html#syntactic-sugar-for-indexed-update-operators>`_.
+  modificatons. For further information, see `Indexed Update Operators
+  <https://jax.readthedocs.io/en/latest/jax.ops.html#indexed-update-operators>`_.
 
   In particular:
 
   - ``x = x.at[idx].set(y)`` is a pure equivalent of ``x[idx] = y``.
   - ``x = x.at[idx].add(y)`` is a pure equivalent of ``x[idx] += y``.
-  - ``x = x.at[idx].mul(y)`` is a pure equivalent of ``x[idx] *= y``.
+  - ``x = x.at[idx].multiply(y)`` (aka ``mul``) is a pure equivalent of
+    ``x[idx] *= y``.
+  - ``x = x.at[idx].divide(y)`` is a pure equivalent of ``x[idx] /= y``.
+  - ``x = x.at[idx].power(y)`` is a pure equivalent of ``x[idx] **= y``.
   - ``x = x.at[idx].min(y)`` is a pure equivalent of
     ``x[idx] = minimum(x[idx], y)``.
   - ``x = x.at[idx].max(y)`` is a pure equivalent of
@@ -5866,6 +5870,8 @@ def __getitem__(self, index):
   def __repr__(self):
     return f"_IndexUpdateHelper({repr(self.array)})"
 
+_power = power
+_divide = divide
 
 class _IndexUpdateRef:
   """Helper object to call indexed update functions for an (advanced) index.
@@ -5886,74 +5892,100 @@ def __repr__(self):
   def set(self, values, indices_are_sorted=False, unique_indices=False):
     """Pure equivalent of ``x[idx] = y``.
 
-    ``x.at[idx].set(y)`` is syntactic sugar for
-    ``jax.ops.index_update(x, jax.ops.index[idx], y)``, and
-    returns the value of ``x`` that would result from the NumPy-style
+    Returns the value of ``x`` that would result from the NumPy-style
     :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] = y``.
 
     See :mod:`jax.ops` for details.
     """
-    return ops.index_update(self.array, self.index, values,
-                            indices_are_sorted=indices_are_sorted,
-                            unique_indices=unique_indices)
+    return scatter._scatter_update(self.array, self.index, values, lax.scatter,
+                                   indices_are_sorted=indices_are_sorted,
+                                   unique_indices=unique_indices)
 
   def add(self, values, indices_are_sorted=False, unique_indices=False):
     """Pure equivalent of ``x[idx] += y``.
 
-    ``x.at[idx].add(y)`` is syntactic sugar for
-    ``jax.ops.index_add(x, jax.ops.index[idx], y)``, and
-    returns the value of ``x`` that would result from the NumPy-style
+    Returns the value of ``x`` that would result from the NumPy-style
     :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] += y``.
 
     See :mod:`jax.ops` for details.
     """
-    return ops.index_add(self.array, self.index, values,
-                         indices_are_sorted=indices_are_sorted,
-                         unique_indices=unique_indices)
+    return scatter._scatter_update(self.array, self.index, values,
+                                   lax.scatter_add,
+                                   indices_are_sorted=indices_are_sorted,
+                                   unique_indices=unique_indices)
 
-  def mul(self, values, indices_are_sorted=False, unique_indices=False):
-    """Pure equivalent of ``x[idx] += y``.
+  def multiply(self, values, indices_are_sorted=False, unique_indices=False):
+    """Pure equivalent of ``x[idx] *= y``.
 
-    ``x.at[idx].mul(y)`` is syntactic sugar for
-    ``jax.ops.index_mul(x, jax.ops.index[idx], y)``, and
-    returns the value of ``x`` that would result from the NumPy-style
+    Returns the value of ``x`` that would result from the NumPy-style
     :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] *= y``.
 
     See :mod:`jax.ops` for details.
     """
-    return ops.index_mul(self.array, self.index, values,
-                         indices_are_sorted=indices_are_sorted,
-                         unique_indices=unique_indices)
+    return scatter._scatter_update(self.array, self.index, values,
+                                   lax.scatter_mul,
+                                   indices_are_sorted=indices_are_sorted,
+                                   unique_indices=unique_indices)
+  mul = multiply
+
+  def divide(self, values, indices_are_sorted=False, unique_indices=False):
+    """Pure equivalent of ``x[idx] /= y``.
+
+    Returns the value of ``x`` that would result from the NumPy-style
+    :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] /= y``.
+
+    See :mod:`jax.ops` for details.
+    """
+    return _divide(
+      self.array,
+      scatter._scatter_update(ones_like(self.array), self.index, values,
+                              lax.scatter_mul,
+                              indices_are_sorted=indices_are_sorted,
+                              unique_indices=unique_indices))
+
+  def power(self, values, indices_are_sorted=False, unique_indices=False):
+    """Pure equivalent of ``x[idx] **= y``.
+
+    Returns the value of ``x`` that would result from the NumPy-style
+    :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] **= y``.
+
+    See :mod:`jax.ops` for details.
+    """
+    return _power(
+      self.array,
+      scatter._scatter_update(ones_like(self.array), self.index, values,
+                              lax.scatter_mul,
+                              indices_are_sorted=indices_are_sorted,
+                              unique_indices=unique_indices))
 
   def min(self, values, indices_are_sorted=False, unique_indices=False):
     """Pure equivalent of ``x[idx] = minimum(x[idx], y)``.
 
-    ``x.at[idx].min(y)`` is syntactic sugar for
-    ``jax.ops.index_min(x, jax.ops.index[idx], y)``, and
-    returns the value of ``x`` that would result from the NumPy-style
+    Returns the value of ``x`` that would result from the NumPy-style
     :mod:indexed assignment <numpy.doc.indexing>`
     ``x[idx] = minimum(x[idx], y)``.
 
     See :mod:`jax.ops` for details.
     """
-    return ops.index_min(self.array, self.index, values,
-                         indices_are_sorted=indices_are_sorted,
-                         unique_indices=unique_indices)
+    return scatter._scatter_update(self.array, self.index, values,
+                                   lax.scatter_min,
+                                   indices_are_sorted=indices_are_sorted,
+                                   unique_indices=unique_indices)
 
   def max(self, values, indices_are_sorted=False, unique_indices=False):
     """Pure equivalent of ``x[idx] = maximum(x[idx], y)``.
 
-    ``x.at[idx].max(y)`` is syntactic sugar for
-    ``jax.ops.index_max(x, jax.ops.index[idx], y)``, and
-    returns the value of ``x`` that would result from the NumPy-style
+    Returns the value of ``x`` that would result from the NumPy-style
     :mod:indexed assignment <numpy.doc.indexing>`
     ``x[idx] = maximum(x[idx], y)``.
 
     See :mod:`jax.ops` for details.
     """
-    return ops.index_max(self.array, self.index, values,
-                         indices_are_sorted=indices_are_sorted,
-                         unique_indices=unique_indices)
+    return scatter._scatter_update(self.array, self.index, values,
+                                   lax.scatter_max,
+                                   indices_are_sorted=indices_are_sorted,
+                                   unique_indices=unique_indices)
+
 
 setattr(_DeviceArray, "at", property(_IndexUpdateHelper))
 setattr(_CppDeviceArray, "at", property(_IndexUpdateHelper))


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -138,7 +138,10 @@ def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=''):
   kw = {}
   if atol: kw["atol"] = atol
   if rtol: kw["rtol"] = rtol
-  np.testing.assert_allclose(a, b, **kw, err_msg=err_msg)
+  with np.errstate(invalid='ignore'):
+    # TODO(phawkins): surprisingly, assert_allclose sometimes reports invalid
+    # value errors. It should not do that.
+    np.testing.assert_allclose(a, b, **kw, err_msg=err_msg)
 
 def tolerance(dtype, tol=None):
   tol = {} if tol is None else tol
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
@@ -849,8 +849,10 @@ class UpdateOps(enum.Enum):
   UPDATE = 0
   ADD = 1
   MUL = 2
-  MIN = 3
-  MAX = 4
+  DIV = 3
+  POW = 4
+  MIN = 5
+  MAX = 6
 
   def np_fn(op, indexer, x, y):
     x = x.copy()
@@ -858,6 +860,10 @@ def np_fn(op, indexer, x, y):
       UpdateOps.UPDATE: lambda: y,
       UpdateOps.ADD: lambda: x[indexer] + y,
       UpdateOps.MUL: lambda: x[indexer] * y,
+      UpdateOps.DIV: jtu.ignore_warning(category=RuntimeWarning)(
+        lambda: x[indexer] / y.astype(x.dtype)),
+      UpdateOps.POW: jtu.ignore_warning(category=RuntimeWarning)(
+        lambda: x[indexer] ** y.astype(x.dtype)),
       UpdateOps.MIN: lambda: np.minimum(x[indexer], y),
       UpdateOps.MAX: lambda: np.maximum(x[indexer], y),
     }[op]()
@@ -880,12 +886,21 @@ def sugar_fn(op, indexer, x, y, indices_are_sorted=False,
     return {
       UpdateOps.UPDATE: x.at[indexer].set,
       UpdateOps.ADD: x.at[indexer].add,
-      UpdateOps.MUL: x.at[indexer].mul,
+      UpdateOps.MUL: x.at[indexer].multiply,
+      UpdateOps.DIV: x.at[indexer].divide,
+      UpdateOps.POW: x.at[indexer].power,
       UpdateOps.MIN: x.at[indexer].min,
       UpdateOps.MAX: x.at[indexer].max,
     }[op](y, indices_are_sorted=indices_are_sorted,
           unique_indices=unique_indices)
 
+  def dtypes(op):
+    if op == UpdateOps.UPDATE:
+      return all_dtypes
+    elif op == UpdateOps.DIV or op == UpdateOps.POW:
+      return jtu.dtypes.inexact
+    else:
+      return default_dtypes
 
 class IndexedUpdateTest(jtu.JaxTestCase):
 
@@ -899,10 +914,10 @@ class IndexedUpdateTest(jtu.JaxTestCase):
   } for name, index_specs in s(STATIC_INDEXING_TESTS)
     for shape, indexer in s(index_specs)
     for op in s(UpdateOps)
-    for dtype in s(all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
+    for dtype in s(UpdateOps.dtypes(op))
     for update_shape in s(_broadcastable_shapes(_update_shape(shape, indexer)))
     for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes)
-    for sugared in s([True, False]))))
+    for sugared in (s([True, False]) if op not in [UpdateOps.DIV, UpdateOps.POW] else [True]))))
   def testStaticIndexing(self, shape, dtype, update_shape, update_dtype,
                          indexer, sugared, op):
     rng = jtu.rand_default(self.rng())
@@ -912,87 +927,78 @@ def testStaticIndexing(self, shape, dtype, update_shape, update_dtype,
       jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y)
     else:
       jax_fn = lambda x, y: UpdateOps.jax_fn(op, indexer, x, y)
-    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker)
+    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker,
+                            tol={np.complex128: 1e-14})
     self._CompileAndCheck(jax_fn, args_maker)
 
   @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
-      "testcase_name": "{}_inshape={}_indexer={}_update={}_sugared={}_op={}".format(
+      "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), sugared, op.name),
+          jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
        "shape": shape, "dtype": dtype, "indexer": indexer,
        "update_shape": update_shape, "update_dtype": update_dtype,
-       "op": op, "sugared": sugared
+       "op": op
   } for name, index_specs in s(ADVANCED_INDEXING_TESTS_NO_REPEATS)
     for shape, indexer in s(index_specs)
     for op in s(UpdateOps)
-    for dtype in s(all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
+    for dtype in s(UpdateOps.dtypes(op))
     for update_shape in s(_broadcastable_shapes(_update_shape(shape, indexer)))
-    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes)
-    for sugared in s([True, False]))))
+    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes))))
   def testAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
-                           indexer, sugared, op):
+                           indexer, op):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
     np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
-    if sugared:
-      jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y, unique_indices=True)
-    else:
-      jax_fn = lambda x, y: UpdateOps.jax_fn(op, indexer, x, y, unique_indices=True)
-    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker)
+    jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y, unique_indices=True)
+    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker,
+                            tol={np.complex128: 1e-14})
     self._CompileAndCheck(jax_fn, args_maker)
 
   @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
-      "testcase_name": "{}_inshape={}_indexer={}_update={}_sugared={}_op={}".format(
+      "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), sugared, op.name),
+          jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
        "shape": shape, "dtype": dtype, "indexer": indexer,
        "update_shape": update_shape, "update_dtype": update_dtype,
-       "op": op, "sugared": sugared
+       "op": op
   } for name, index_specs in s(ADVANCED_INDEXING_TESTS_NO_REPEATS_SORTED)
     for shape, indexer in s(index_specs)
     for op in s(UpdateOps)
-    for dtype in s(all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
+    for dtype in s(UpdateOps.dtypes(op))
     for update_shape in s(_broadcastable_shapes(_update_shape(shape, indexer)))
-    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes)
-    for sugared in s([True, False]))))
+    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes))))
   def testAdvancedIndexingSorted(self, shape, dtype, update_shape, update_dtype,
-                           indexer, sugared, op):
+                           indexer, op):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
     np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
-    if sugared:
-      jax_fn = lambda x, y: UpdateOps.sugar_fn(
-          op, indexer, x, y, indices_are_sorted=True, unique_indices=True)
-    else:
-      jax_fn = lambda x, y: UpdateOps.jax_fn(
-          op, indexer, x, y, indices_are_sorted=True, unique_indices=True)
-    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker, check_dtypes=True)
+    jax_fn = lambda x, y: UpdateOps.sugar_fn(
+      op, indexer, x, y, indices_are_sorted=True, unique_indices=True)
+    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker, check_dtypes=True,
+                            tol={np.complex128: 1e-14})
     self._CompileAndCheck(jax_fn, args_maker, check_dtypes=True)
 
   @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
-      "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}_sugared={}".format(
+      "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, sugared),
+          jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
        "shape": shape, "dtype": dtype, "indexer": indexer,
        "update_shape": update_shape, "update_dtype": update_dtype,
-       "op": op, "sugared": sugared
+       "op": op
   } for name, index_specs in s(MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS)
     for shape, indexer in s(index_specs)
     for op in s(UpdateOps)
-    for dtype in s(all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
+    for dtype in s(UpdateOps.dtypes(op))
     for update_shape in s(_broadcastable_shapes(_update_shape(shape, indexer)))
-    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes)
-    for sugared in s([True, False]))))
+    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes))))
   def testMixedAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
-                                indexer, sugared, op):
+                                indexer, op):
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
     np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
-    if sugared:
-      jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y)
-    else:
-      jax_fn = lambda x, y: UpdateOps.jax_fn(op, indexer, x, y)
-    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker)
+    jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y)
+    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker,
+                            tol={np.complex128: 1e-14})
     self._CompileAndCheck(jax_fn, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list({
@@ -1012,7 +1018,7 @@ def testMixedAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
   def testStaticIndexingGrads(self, shape, dtype, update_shape, update_dtype,
                               indexer, op):
     rng = jtu.rand_default(self.rng())
-    jax_fn = lambda x, y: UpdateOps.jax_fn(op, indexer, x, y)
+    jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y)
     x = rng(shape, dtype)
     y = rng(update_shape, update_dtype)
     check_grads(jax_fn, (x, y), 2, rtol=1e-3, atol=1e-3, eps=1.)