Repository: google/jax

Issue:
Update jax.ops.index_* documentation to include current preferred index notation
If you search on Google for how to update array indices in JAX, you end up on [this documentation page](https://jax.readthedocs.io/en/latest/_autosummary/jax.ops.index_update.html#jax-ops-index-update) for `jax.ops.index_update`.

`jax.ops.index_update` is no longer the best way to update indices -- in most cases users should instead write code like `x.at[idx].set(y)`, which is a really great improvement!

The documentation for the `jax.ops.index_*` should also describe and/or link to the alternate, simpler, and preferred notation for working with indices.


Discussion:
Thanks for the report – agreed this is something that should be addressed. Are you interested in contributing a fix? If not I'm happy to take care of it.
I am not going to have time to submit a fix in a reasonable timeframe. If you are willing, you should definitely do it! Thank you!
OK, will do. Thanks!

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
@@ -6690,22 +6690,43 @@ 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 `Indexed Update Operators
-  <https://jax.readthedocs.io/en/latest/jax.ops.html#indexed-update-operators>`_.
+  modificatons.
 
   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].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
-    ``x[idx] = maximum(x[idx], y)``.
+  ==============================  ================================
+  Alternate syntax                Equivalent In-place expression
+  ==============================  ================================
+  ``x = x.at[idx].set(y)``        ``x[idx] = y``
+  ``x = x.at[idx].add(y)``        ``x[idx] += y``
+  ``x = x.at[idx].multiply(y)``   ``x[idx] *= y``
+  ``x = x.at[idx].divide(y)``     ``x[idx] /= y``
+  ``x = x.at[idx].power(y)``      ``x[idx] **= y``
+  ``x = x.at[idx].min(y)``        ``x[idx] = minimum(x[idx], y)``
+  ``x = x.at[idx].max(y)``        ``x[idx] = maximum(x[idx], y)``
+  ``x = x.at[idx].get()``         ``x = x[idx]``
+  ==============================  ================================
+
+  None of these expressions modify the original ``x``; instead they return
+  a modified copy of ``x``. However, inside a :py:func:`jax.jit` compiled function,
+  expressions like ``x = x.at[idx].set(y)`` are guaranteed to be applied in-place.
+
+  Unlike NumPy in-place operations such as :code:`x[idx] += y`, if multiple
+  indices refer to the same location, all updates will be applied (NumPy would
+  only apply the last update, rather than applying all updates.) The order
+  in which conflicting updates are applied is implementation-defined and may be
+  nondeterministic (e.g., due to concurrency on some hardware platforms).
+
+  By default, JAX assumes that all indices are in-bounds. There is experimental
+  support for giving more precise semantics to out-of-bounds indexed accesses,
+  via the ``mode`` parameter to functions such as ``get`` and ``set``. Valid
+  values for ``mode`` include ``"clip"``, which means that out-of-bounds indices
+  will be clamped into range, and ``"fill"``/``"drop"``, which are aliases and
+  mean that out-of-bounds reads will be filled with a scalar ``fill_value``,
+  and out-of-bounds writes will be discarded.
   """
+  # TODO(jakevdp): document additional arguments to the methods, including
+  # `indices_are_sorted`, `unique_indices`, `mode`, and `fill_value`.
   __slots__ = ("array",)
 
   def __init__(self, array):
@@ -6716,6 +6737,7 @@ def __getitem__(self, index):
 
   def __repr__(self):
     return f"_IndexUpdateHelper({repr(self.array)})"
+ndarray.at.__doc__ = _IndexUpdateHelper.__doc__
 
 _power_fn = power
 _divide_fn = divide
diff --git a/jax/_src/ops/scatter.py b/jax/_src/ops/scatter.py
--- a/jax/_src/ops/scatter.py
+++ b/jax/_src/ops/scatter.py
@@ -114,6 +114,9 @@ def _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx,
 class _Indexable(object):
   """Helper object for building indexes for indexed update functions.
 
+  .. deprecated:: 0.2.22
+     Prefer the use of :attr:`jax.numpy.ndarray.at`.
+
   This is a singleton object that overrides the :code:`__getitem__` method
   to return the index it is passed.
 
@@ -136,6 +139,9 @@ def index_add(x: Array,
               unique_indices: bool = False) -> Array:
   """Pure equivalent of :code:`x[idx] += y`.
 
+  .. deprecated:: 0.2.22
+     Prefer the use of :attr:`jax.numpy.ndarray.at`.
+
   Returns the value of `x` that would result from the
   NumPy-style :mod:`indexed assignment <numpy.doc.indexing>`::
 
@@ -183,6 +189,9 @@ def index_mul(x: Array,
               unique_indices: bool = False) -> Array:
   """Pure equivalent of :code:`x[idx] *= y`.
 
+  .. deprecated:: 0.2.22
+     Prefer the use of :attr:`jax.numpy.ndarray.at`.
+
   Returns the value of `x` that would result from the
   NumPy-style :mod:`indexed assignment <numpy.doc.indexing>`::
 
@@ -230,6 +239,9 @@ def index_min(x: Array,
               unique_indices: bool = False) -> Array:
   """Pure equivalent of :code:`x[idx] = minimum(x[idx], y)`.
 
+  .. deprecated:: 0.2.22
+     Prefer the use of :attr:`jax.numpy.ndarray.at`.
+
   Returns the value of `x` that would result from the
   NumPy-style :mod:`indexed assignment <numpy.doc.indexing>`::
 
@@ -274,6 +286,9 @@ def index_max(x: Array,
               unique_indices: bool = False) -> Array:
   """Pure equivalent of :code:`x[idx] = maximum(x[idx], y)`.
 
+  .. deprecated:: 0.2.22
+     Prefer the use of :attr:`jax.numpy.ndarray.at`.
+
   Returns the value of `x` that would result from the
   NumPy-style :mod:`indexed assignment <numpy.doc.indexing>`::
 
@@ -318,6 +333,9 @@ def index_update(x: Array,
                  unique_indices: bool = False) -> Array:
   """Pure equivalent of :code:`x[idx] = y`.
 
+  .. deprecated:: 0.2.22
+     Prefer the use of :attr:`jax.numpy.ndarray.at`.
+
   Returns the value of `x` that would result from the
   NumPy-style :mod:`indexed assignment <numpy.doc.indexing>`::
 


Test patch: