corrected jax expression, numpy function

keras-team · Sep 21, 2023 · 5efabf5 · 5efabf5
1 parent 7f76459
commit 5efabf5
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/keras_core/backend/jax/math.py b/keras_core/backend/jax/math.py
@@ -251,4 +251,4 @@ def rsqrt(x):
 
 
 def erf(x):
-    return jnp.erf(x)
+    return jax.lax.erf(x)
diff --git a/keras_core/backend/numpy/math.py b/keras_core/backend/numpy/math.py
@@ -305,4 +305,4 @@ def rsqrt(x):
 
 
 def erf(x):
-    return scipy.special.erf(x)
+    return np.array(scipy.special.erf(x))
diff --git a/keras_core/ops/math_test.py b/keras_core/ops/math_test.py
@@ -846,7 +846,7 @@ def test_erf_operation_basic(self):
         )
 
         # Output from the erf operation in keras_core
-        output_from_erf_op = kmath.erf(sample_values).numpy()
+        output_from_erf_op = kmath.erf(sample_values)
 
         # Assert that the outputs are close
         self.assertAllClose(expected_output, output_from_erf_op, atol=1e-5)
@@ -860,7 +860,7 @@ def test_erf_operation_dtype(self):
             expected_output = (2 / np.sqrt(np.pi)) * np.vectorize(math.erf)(
                 sample_values
             )
-            output_from_erf_op = kmath.erf(sample_values).numpy()
+            output_from_erf_op = kmath.erf(sample_values)
             self.assertAllClose(expected_output, output_from_erf_op, atol=1e-5)
 
     def test_erf_operation_edge_cases(self):
@@ -869,7 +869,7 @@ def test_erf_operation_edge_cases(self):
         expected_edge_output = (2 / np.sqrt(np.pi)) * np.vectorize(math.erf)(
             edge_values
         )
-        output_from_edge_erf_op = kmath.erf(edge_values).numpy()
+        output_from_edge_erf_op = kmath.erf(edge_values)
         self.assertAllClose(
             expected_edge_output, output_from_edge_erf_op, atol=1e-5
         )