From dee0b1dc9f3ea7804b6a40d97ddfc288610beb9d Mon Sep 17 00:00:00 2001 From: samthakur587 Date: Sun, 31 Dec 2023 19:41:14 +0530 Subject: [PATCH] fix: fix failing test case at jax backend --- ivy/functional/backends/jax/activations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ivy/functional/backends/jax/activations.py b/ivy/functional/backends/jax/activations.py index 08f4f42c88ebf..cdb9340436cf6 100644 --- a/ivy/functional/backends/jax/activations.py +++ b/ivy/functional/backends/jax/activations.py @@ -37,7 +37,7 @@ def leaky_relu( def relu( x: JaxArray, /, *, complex_mode="jax", out: Optional[JaxArray] = None ) -> JaxArray: - return jnp.maximum(x, 0) + return jax.nn.relu(x) def sigmoid(