Repository: google/jax

Issue:
Keyword arguments to jnp.ndarray.clip() do not follow the documentation
The following should work according to the docs (https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.html), and it does work when using `numpy` instead of `jax.numpy`:

```python
import jax.numpy as jnp
jnp.zeros(3).clip(min=1)
```

However, in contrast to the docs, `jnp.ndarray.clip` actually expects the arguments to be called `a_min` and `a_max` respectively, just like `jnp.clip` does – i.e. `jnp.zeros(3).clip(a_min=1)` works.


Discussion:
Thanks for the report – currently the DeviceArray `clip` method delegates directly to `jnp.clip`:
https://github.com/google/jax/blob/0dee3550250cf8e0afe8f1c0d24399844648584f/jax/_src/numpy/lax_numpy.py#L6085
https://github.com/google/jax/blob/0dee3550250cf8e0afe8f1c0d24399844648584f/jax/_src/numpy/lax_numpy.py#L6332-L6333

It looks like that has to be changed, given the different API between the function and the method. Are you interested in sending a PR?

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
@@ -5946,6 +5946,20 @@ def _nbytes(arr):
   return size(arr) * _dtype(arr).itemsize
 
 
+def _clip(number, min=None, max=None, out=None, *, a_min=None, a_max=None):
+  # ndarray.clip has a slightly different API from clip (min -> a_min, max -> a_max)
+  # TODO: remove after deprecation window
+  if a_min is not None or a_max is not None:
+    warnings.warn('`a_min` and `a_max` keyword arguments to ndarray.clip are deprecated '
+                  'in favor of `min` and `max` for compatibility with numpy. '
+                  'They will be removed in JAX 0.22.2', FutureWarning)
+  if min is None and a_min is not None:
+    min = a_min
+  if max is None and a_max is not None:
+    max = a_max
+  return clip(number, a_min=min, a_max=max, out=out)
+
+
 def _view(arr, dtype=None, type=None):
   lax._check_user_dtype_supported(dtype, "view")
   if type is not None:
@@ -6082,7 +6096,7 @@ def _operator_round(number, ndigits=None):
 # These numpy.ndarray methods are just refs to an equivalent numpy function
 _nondiff_methods = ["all", "any", "argmax", "argmin", "argpartition", "argsort",
                     "nonzero", "searchsorted", "round"]
-_diff_methods = ["choose", "clip", "conj", "conjugate", "cumprod", "cumsum",
+_diff_methods = ["choose", "conj", "conjugate", "cumprod", "cumsum",
                  "diagonal", "dot", "max", "mean", "min", "prod", "ptp",
                  "ravel", "repeat", "sort", "squeeze", "std", "sum",
                  "swapaxes", "take", "tile", "trace", "var"]
@@ -6313,6 +6327,7 @@ def _set_shaped_array_attributes(shaped_array):
   setattr(shaped_array, "astype", core.aval_method(_astype))
   setattr(shaped_array, "view", core.aval_method(_view))
   setattr(shaped_array, "nbytes", core.aval_property(_nbytes))
+  setattr(shaped_array, "clip", core.aval_method(_clip))
 
   setattr(shaped_array, "_array_module", staticmethod(__array_module__))
   setattr(shaped_array, "broadcast", core.aval_method(lax.broadcast))
@@ -6340,6 +6355,7 @@ def _set_device_array_base_attributes(device_array):
   setattr(device_array, "astype", _astype)
   setattr(device_array, "view", _view)
   setattr(device_array, "nbytes", property(_nbytes))
+  setattr(device_array, "clip", _clip)
 
 _set_device_array_base_attributes(DeviceArray)
 


Test patch: