Dispatch Clip using jnp.clip

Author: rlouf

aesara/link/jax/dispatch/scalar.py

@@ -55,8 +55,8 @@ def identity(x):
 
 @jax_funcify.register(Clip)
 def jax_funcify_Clip(op, **kwargs):
-    def clip(x, min, max):
-        return jnp.where(x < min, min, jnp.where(x > max, max, x))
+    def clip(x, a_min, a_max):
+        return jnp.clip(x, a_min, a_max)
 
     return clip