Repository: google/jax

Issue:
Make arguments of index update helper functions keyword only
For instance, the following code will fail silently, especially given that `at` is not type hinted. One could argue it's just the user (me) being silly, but it's still surprising and difficult to debug. Making arguments like `mode` keyword-only will eliminate this kind of problem.

```python
I = jnp.array([0, 10])
a = jnp.arange(5).at[I].get('fill', fill_value=-1)
#     this is actually: get(indices_are_sorted='fill', fill_value=-1)
```



Discussion:
This is a great idea – let me explore a bit to see how many downstream projects it would break.

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
@@ -26,7 +26,8 @@
 
 import builtins
 import collections
-from functools import partial
+from functools import partial, wraps
+import inspect
 import math
 import operator
 import types
@@ -5326,6 +5327,27 @@ def __repr__(self):
 Array.at.__doc__ = _IndexUpdateHelper.__doc__
 
 
+# TODO(jakevdp): remove these deprecation warnings after June 2023
+def allow_pass_by_position_with_warning(f):
+  @wraps(f)
+  def wrapped(*args, **kwargs):
+    sig = inspect.signature(f)
+    try:
+      sig.bind(*args, **kwargs)
+    except TypeError:
+      argspec = inspect.getfullargspec(f)
+      n_positional = len(argspec.args)
+      keywords = argspec.kwonlyargs[:len(args) - n_positional]
+      warnings.warn(
+          f"jnp.ndarray.at[...].{f.__name__}: Passing '{keywords[0]}' by position is deprecated. "
+          f"Pass by keyword instead", category=FutureWarning, stacklevel=2)
+      converted_kwargs = dict(zip(keywords, args[n_positional:]))
+      return f(*args[:n_positional], **converted_kwargs, **kwargs)
+    else:
+      return f(*args, **kwargs)
+  return wrapped
+
+
 class _IndexUpdateRef:
   """Helper object to call indexed update functions for an (advanced) index.
 
@@ -5342,7 +5364,8 @@ def __init__(self, array, index):
   def __repr__(self):
     return f"_IndexUpdateRef({repr(self.array)}, {repr(self.index)})"
 
-  def get(self, indices_are_sorted=False, unique_indices=False,
+  @allow_pass_by_position_with_warning
+  def get(self, *, indices_are_sorted=False, unique_indices=False,
           mode=None, fill_value=None):
     """Equivalent to ``x[idx]``.
 
@@ -5358,7 +5381,8 @@ def get(self, indices_are_sorted=False, unique_indices=False,
                            unique_indices=unique_indices, mode=mode,
                            fill_value=fill_value)
 
-  def set(self, values, indices_are_sorted=False, unique_indices=False,
+  @allow_pass_by_position_with_warning
+  def set(self, values, *, indices_are_sorted=False, unique_indices=False,
           mode=None):
     """Pure equivalent of ``x[idx] = y``.
 
@@ -5371,7 +5395,8 @@ def set(self, values, indices_are_sorted=False, unique_indices=False,
                                    indices_are_sorted=indices_are_sorted,
                                    unique_indices=unique_indices, mode=mode)
 
-  def apply(self, func, indices_are_sorted=False, unique_indices=False,
+  @allow_pass_by_position_with_warning
+  def apply(self, func, *, indices_are_sorted=False, unique_indices=False,
             mode=None):
     """Pure equivalent of ``func.at(x, idx)`` for a unary ufunc ``func``.
 
@@ -5394,7 +5419,8 @@ def _scatter_apply(x, indices, _, dims, **kwargs):
                                    indices_are_sorted=indices_are_sorted,
                                    unique_indices=unique_indices, mode=mode)
 
-  def add(self, values, indices_are_sorted=False, unique_indices=False,
+  @allow_pass_by_position_with_warning
+  def add(self, values, *, indices_are_sorted=False, unique_indices=False,
           mode=None):
     """Pure equivalent of ``x[idx] += y``.
 
@@ -5408,7 +5434,8 @@ def add(self, values, indices_are_sorted=False, unique_indices=False,
                                    indices_are_sorted=indices_are_sorted,
                                    unique_indices=unique_indices, mode=mode)
 
-  def multiply(self, values, indices_are_sorted=False, unique_indices=False,
+  @allow_pass_by_position_with_warning
+  def multiply(self, values, *, indices_are_sorted=False, unique_indices=False,
                mode=None):
     """Pure equivalent of ``x[idx] *= y``.
 
@@ -5424,7 +5451,8 @@ def multiply(self, values, indices_are_sorted=False, unique_indices=False,
                                    mode=mode)
   mul = multiply
 
-  def divide(self, values, indices_are_sorted=False, unique_indices=False,
+  @allow_pass_by_position_with_warning
+  def divide(self, values, *, indices_are_sorted=False, unique_indices=False,
              mode=None):
     """Pure equivalent of ``x[idx] /= y``.
 
@@ -5440,7 +5468,8 @@ def divide(self, values, indices_are_sorted=False, unique_indices=False,
                               indices_are_sorted=indices_are_sorted,
                               unique_indices=unique_indices, mode=mode))
 
-  def power(self, values, indices_are_sorted=False, unique_indices=False,
+  @allow_pass_by_position_with_warning
+  def power(self, values, *, indices_are_sorted=False, unique_indices=False,
             mode=None):
     """Pure equivalent of ``x[idx] **= y``.
 
@@ -5456,7 +5485,8 @@ def power(self, values, indices_are_sorted=False, unique_indices=False,
                               indices_are_sorted=indices_are_sorted,
                               unique_indices=unique_indices, mode=mode))
 
-  def min(self, values, indices_are_sorted=False, unique_indices=False,  # noqa: F811
+  @allow_pass_by_position_with_warning
+  def min(self, values, *, indices_are_sorted=False, unique_indices=False,
           mode=None):
     """Pure equivalent of ``x[idx] = minimum(x[idx], y)``.
 
@@ -5471,7 +5501,8 @@ def min(self, values, indices_are_sorted=False, unique_indices=False,  # noqa: F
                                    indices_are_sorted=indices_are_sorted,
                                    unique_indices=unique_indices, mode=mode)
 
-  def max(self, values, indices_are_sorted=False, unique_indices=False,  # noqa: F811
+  @allow_pass_by_position_with_warning
+  def max(self, values, *, indices_are_sorted=False, unique_indices=False,
           mode=None):
     """Pure equivalent of ``x[idx] = maximum(x[idx], y)``.
 


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
@@ -929,6 +929,13 @@ def testFloatIndexingError(self):
     with self.assertRaisesRegex(TypeError, BAD_INDEX_TYPE_ERROR):
       jnp.zeros(2).at[0.].set(1.)
 
+  def testIndexingPositionalArgumentWarning(self):
+    x = jnp.arange(4)
+    with self.assertWarnsRegex(
+        FutureWarning,  "Passing 'indices_are_sorted' by position is deprecated"):
+      out = x.at[5].set(1, True, mode='drop')
+    self.assertArraysEqual(out, x)
+
   def testIndexOutOfBounds(self):  # https://github.com/google/jax/issues/2245
     x = jnp.arange(5, dtype=jnp.int32) + 1
     self.assertAllClose(x, x[:10])