Repository: google/jax

Issue:
Make it easier to find documentation for the DeviceArray.at attribute
It's currently very hard to find documentation about the syntax for slice updates. I'm not sure if I'm unable to find something on it because "at" is a stop word, or if isn't documented at all. I only was able to find out about this feature because of the FLAX documentation.


Discussion:
Hi - thanks for opening the issue. The attribute is documented here: https://jax.readthedocs.io/en/latest/jax.ops.html#syntactic-sugar-for-indexed-update-operators

I'm not sure how to make this more discoverable. Do you have any suggestions?
Hi @jakevdp 

> Do you have any suggestions?

In the Glossary https://jax.readthedocs.io/en/latest/glossary.html where it says:

```
`DeviceArray`

  JAX’s analog of the `numpy.ndarray`. See `jax.interpreters.xla.DeviceArray`.
```

`jax.interpreters.xla.DeviceArray` is not "clickable". Maybe fixing that could resolve the discoverability for some users?
And then, you could maybe add a reference to https://jax.readthedocs.io/en/latest/jax.ops.html#syntactic-sugar-for-indexed-update-operators in the `jax.interpreters.xla.DeviceArray` API docs because the two can be (f)used together.
Ah, yeah that's a good idea. It's a bit non-trivial, though, because `DeviceArray` doesn't actually have the `at` attribute: instead it is defined on the two implementations of `DeviceArray` that are used by the Python JIT and the C++ JIT respectively. Offhand I don't know any way to make sphinx document that in a non-confusing way (this is why it's not been done already). I'll try a few things and see if I can figure out a workaround.
> Hi @jakevdp
> 
> > Do you have any suggestions?
> 
> In the Glossary https://jax.readthedocs.io/en/latest/glossary.html where it says:
> 
> ```
> `DeviceArray`
> 
>   JAX’s analog of the `numpy.ndarray`. See `jax.interpreters.xla.DeviceArray`.
> ```
> 
> `jax.interpreters.xla.DeviceArray` is not "clickable". Maybe fixing that could resolve the discoverability for some users?

This is basically exactly why I had trouble finding it I think. There being two implementations makes sense, but having some docs on their shared interface sounds great.

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
@@ -5837,18 +5837,20 @@ class _IndexUpdateHelper:
   # Note: this docstring will appear as the docstring for the `at` property.
   """Indexable helper object to call indexed update functions.
 
-  The `at` property is syntactic sugar for calling the indexed update functions
+  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.
+  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>`_.
 
   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].min(y)`` is a pure equivalent of
-      ``x[idx] = minimum(x[idx], y)``.
+    ``x[idx] = minimum(x[idx], y)``.
   - ``x = x.at[idx].max(y)`` is a pure equivalent of
-      ``x[idx] = maximum(x[idx], y)``.
+    ``x[idx] = maximum(x[idx], y)``.
   """
   __slots__ = ("array",)
 


Test patch: