Merge pull request #29001 from johannahaffner:test-clip

PiperOrigin-RevId: 763865376

Author: Google-ML-Automation

jax/_src/numpy/lax_numpy.py

@@ -3410,6 +3410,7 @@ def clip(
   Returns:
     An array containing values from ``arr``, with values smaller than ``min`` set
     to ``min``, and values larger than ``max`` set to ``max``.
+    Wherever ``min`` is larger than ``max``, the value of ``max`` is returned.
 
   See also:
     - :func:`jax.numpy.minimum`: Compute the element-wise minimum value of two arrays.

tests/lax_numpy_test.py

@@ -1065,6 +1065,14 @@ def testClipDeprecatedArgs(self):
                                              "Passing arguments 'a', 'a_min' or 'a_max' to jax.numpy.clip is deprecated"):
       jnp.clip(jnp.arange(4), a_min=2, a_max=3)
 
+  def testClipUpperPrecedence(self):
+    a_min = 3 * np.ones(1)
+    a_max = 2 * np.ones(1)
+    x = 4 * np.ones(1)
+    y = jnp.clip(x, min=a_min, max=a_max)
+    assert y == a_max, f"Expected {y} to equal {a_max} when a_min > a_max."
+    assert y == jnp.asarray(np.clip(x, a_min=a_min, a_max=a_max))
+
   def testHypotComplexInputError(self):
     rng = jtu.rand_default(self.rng())
     x = rng((5,), dtype=jnp.complex64)