code reformatting done, introduced erf function in jax, torch and numpy

keras-team · sqali · Sep 17, 2023 · Sep 18, 2023 · Sep 18, 2023 · Sep 18, 2023
commit 61c9690a1b91d5e70927b2b9862cd5525b5fd1b6
diff --git a/keras_core/backend/jax/math.py b/keras_core/backend/jax/math.py
@@ -248,3 +248,7 @@ def istft(
 
 def rsqrt(x):
     return jax.lax.rsqrt(x)
+
+
+def erf(x):
+    return jnp.erf(x)
diff --git a/keras_core/backend/numpy/math.py b/keras_core/backend/numpy/math.py
@@ -302,3 +302,7 @@ def istft(
 
 def rsqrt(x):
     return 1.0 / np.sqrt(x)
+
+
+def erf(x):
+    return scipy.special.erf(x)
diff --git a/keras_core/backend/tensorflow/math.py b/keras_core/backend/tensorflow/math.py
@@ -242,4 +242,4 @@ def rsqrt(x):
 
 
 def erf(x):
-    return tf.math.erf(x)
+    return tf.math.erf(x)
diff --git a/keras_core/backend/torch/math.py b/keras_core/backend/torch/math.py
@@ -408,3 +408,9 @@ def istft(
 def rsqrt(x):
     x = convert_to_tensor(x)
     return torch.rsqrt(x)
+
+
+def erf(x):
+    if not isinstance(x, torch.Tensor):
+        x = torch.tensor(x)
+    return torch.erf(x)
diff --git a/keras_core/ops/math.py b/keras_core/ops/math.py
@@ -936,7 +936,7 @@ class Erf(Operation):
 
     Args:
         input_tensor: A tensor of type `float32` or `float64`.
-    
+
     Returns:
         A tensor of the same shape and type as `input_tensor`, containing the error function values.
 
@@ -952,6 +952,7 @@ class Erf(Operation):
     >>> x_large = np.array([1e10, -1e10])
     >>> y_large = Erf()(x_large)
     """
+
     def __init__(self):
         super().__init__()
 
@@ -961,6 +962,7 @@ def compute_output_spec(self, input_tensor):
     def call(self, input_tensor):
         return backend.erf(input_tensor)
 
+
 def erf(x):
     """Functional interface to the `Erf` operation."""
-    return Erf()(x)
+    return Erf()(x)
diff --git a/keras_core/ops/math_test.py b/keras_core/ops/math_test.py
@@ -838,13 +838,14 @@ def test_rsqrt(self):
 
 
 class MathOpsCorrectnessTest(testing.TestCase, parameterized.TestCase):
-
     def test_erf_operation_basic(self):
         # Sample values for testing
         sample_values = np.array([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0])
 
         # Expected output using numpy's approximation of the error function
-        expected_output = (2 / np.sqrt(np.pi)) * np.vectorize(math.erf)(sample_values)
+        expected_output = (2 / np.sqrt(np.pi)) * np.vectorize(math.erf)(
+            sample_values
+        )
 
         # Output from the erf operation in keras_core
         output_from_erf_op = kmath.erf(sample_values).numpy()
@@ -855,14 +856,22 @@ def test_erf_operation_basic(self):
     def test_erf_operation_dtype(self):
         # Test for float32 and float64 data types
         for dtype in [np.float32, np.float64]:
-            sample_values = np.array([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], dtype=dtype)
-            expected_output = (2 / np.sqrt(np.pi)) * np.vectorize(math.erf)(sample_values)
+            sample_values = np.array(
+                [-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0], dtype=dtype
+            )
+            expected_output = (2 / np.sqrt(np.pi)) * np.vectorize(math.erf)(
+                sample_values
+            )
             output_from_erf_op = kmath.erf(sample_values).numpy()
             self.assertAllClose(expected_output, output_from_erf_op, atol=1e-5)
 
     def test_erf_operation_edge_cases(self):
         # Test for edge cases
         edge_values = np.array([1e10, -1e10, 1e-10, -1e-10], dtype=np.float64)
-        expected_edge_output = (2 / np.sqrt(np.pi)) * np.vectorize(math.erf)(edge_values)
+        expected_edge_output = (2 / np.sqrt(np.pi)) * np.vectorize(math.erf)(
+            edge_values
+        )
         output_from_edge_erf_op = kmath.erf(edge_values).numpy()
-        self.assertAllClose(expected_edge_output, output_from_edge_erf_op, atol=1e-5)
+        self.assertAllClose(
+            expected_edge_output, output_from_edge_erf_op, atol=1e-5
+        )
Original file line number	Diff line number	Diff line change
Expand Up		@@ -242,4 +242,4 @@ def rsqrt(x):


		def erf(x):
		return tf.math.erf(x)
		return tf.math.erf(x)