Repository: google/jax

Issue:
jax.numpy.clip has unexpected behavior which diverges from numpy.clip 
Hi all!

I've recently encountered an issue with `jax.numpy.clip` which produces variable outputs based on the `dtype` of the input (I have jax version `0.2.0`). For example, if we define our `a_min` and `a_max` as floats, but pass an integer type for the `a` parameter, then the limits get cast to the same integer type as the input (see [here](https://github.com/google/jax/blob/ebc5e8bfd6d091bf1fbf67de1565705f30159aaf/jax/numpy/lax_numpy.py#L1514) and [here](https://github.com/google/jax/blob/ebc5e8bfd6d091bf1fbf67de1565705f30159aaf/jax/numpy/lax_numpy.py#L1518)). This can have some undesired results, especially since users might be accustomed to the `numpy.clip` behavior which does the exact opposite. Moreover, this divergence between the two implementations is [not explicitly documented](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.clip.html), as the jax documentation just copies [the numpy doc](https://numpy.org/doc/stable/reference/generated/numpy.clip.html). To detail the issue, I've attached two code snippets:

```python
>>> import jax.numpy as jnp
>>> import numpy as np
>>>
>>> eps = .0001
>>> f = lambda x: jnp.clip(x, eps - 1, 1 - eps)  # x \in [-0.9999, 0.9999] 
>>> g = lambda x: np.clip(x, eps - 1, 1 - eps)   # x \in [-0.9999, 0.9999]
>>>
>>> x = jnp.arange(-3, 3)
>>> f(x)
DeviceArray([0, 0, 0, 0, 0, 0], dtype=int32)
>>> g(x)
array([-0.9999, -0.9999, -0.9999,  0.    ,  0.9999,  0.9999])
```

Considering the implementation, the behavior is as expected only when `jax` can infer the type or the type is explicitly specified:

```python
>>> x = jnp.arange(-3, 3, dtype=jnp.float32)  # Something like jnp.arange(-3.0, 3.0) will also work
>>> f(x)
DeviceArray([-0.9999, -0.9999, -0.9999,  0.    ,  0.9999,  0.9999], dtype=float32)
>>> g(x)
array([-0.9999, -0.9999, -0.9999,  0.    ,  0.9999,  0.9999],
      dtype=float32)
```

My two (mutually exclusive) suggestions would be the following:

* Explicitly document the difference between the two implementations
* Change the jax implementation such that it aligns with the numpy one

I would personally prefer the latter, since not everyone will read the docs, especially if they're already acquainted with the numpy method in the first place, however, I'm not sure what implications this will have on the jax codebase.

Thanks!


Discussion:
Thanks - the issue is here: https://github.com/google/jax/blob/d551cec6e80bc7c31423739e67256c98b453e530/jax/numpy/lax_numpy.py#L1514

JAX coerces the bounds to the dtype of the input array, whereas numpy promotes the array to the dtype of the bounds. I think we should fix this to match numpy's behavior.

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -1510,12 +1510,8 @@ def clip(a, a_min=None, a_max=None):
   if a_min is None and a_max is None:
     raise ValueError("At most one of a_min and a_max may be None")
   if a_min is not None:
-    if _dtype(a_min) != _dtype(a):
-      a_min = lax.convert_element_type(a_min, _dtype(a))
     a = maximum(a_min, a)
   if a_max is not None:
-    if _dtype(a_max) != _dtype(a):
-      a_max = lax.convert_element_type(a_max, _dtype(a))
     a = minimum(a_max, a)
   return a
 


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1081,10 +1081,10 @@ def np_fun(lhs, rhs):
        "shape": shape, "dtype": dtype, "a_min": a_min, "a_max": a_max,
        "rng_factory": jtu.rand_default}
       for shape in all_shapes for dtype in number_dtypes
-      for a_min, a_max in [(-1, None), (None, 1), (-1, 1),
+      for a_min, a_max in [(-1, None), (None, 1), (-0.9, 1),
                            (-np.ones(1), None),
                            (None, np.ones(1)),
-                           (-np.ones(1), np.ones(1))]))
+                           (np.full(1, -0.9), np.ones(1))]))
   def testClipStaticBounds(self, shape, dtype, a_min, a_max, rng_factory):
     rng = rng_factory(self.rng())
     np_fun = lambda x: np.clip(x, a_min=a_min, a_max=a_max)