Implement heaviside function in keras.ops (#21474)

shashaka · web-flow · commit 503bcf56180d · 2025-07-15T08:54:46.000-07:00
* Add heaviside for numpy.py

* Add heaviside for numpy.py

* Add heaviside for ops

* Add test case

* Add test cases

* Correct code by gemini assist

* Correct code by gemini assist

* update excluded_concrete_tests.txt

* Update description to use backticks
diff --git a/keras/api/_tf_keras/keras/ops/__init__.py b/keras/api/_tf_keras/keras/ops/__init__.py
@@ -191,6 +191,7 @@
 from keras.src.ops.numpy import greater_equal as greater_equal
 from keras.src.ops.numpy import hamming as hamming
 from keras.src.ops.numpy import hanning as hanning
+from keras.src.ops.numpy import heaviside as heaviside
 from keras.src.ops.numpy import histogram as histogram
 from keras.src.ops.numpy import hstack as hstack
 from keras.src.ops.numpy import identity as identity
diff --git a/keras/api/_tf_keras/keras/ops/numpy/__init__.py b/keras/api/_tf_keras/keras/ops/numpy/__init__.py
@@ -80,6 +80,7 @@
 from keras.src.ops.numpy import greater_equal as greater_equal
 from keras.src.ops.numpy import hamming as hamming
 from keras.src.ops.numpy import hanning as hanning
+from keras.src.ops.numpy import heaviside as heaviside
 from keras.src.ops.numpy import histogram as histogram
 from keras.src.ops.numpy import hstack as hstack
 from keras.src.ops.numpy import identity as identity
diff --git a/keras/api/ops/__init__.py b/keras/api/ops/__init__.py
@@ -191,6 +191,7 @@
 from keras.src.ops.numpy import greater_equal as greater_equal
 from keras.src.ops.numpy import hamming as hamming
 from keras.src.ops.numpy import hanning as hanning
+from keras.src.ops.numpy import heaviside as heaviside
 from keras.src.ops.numpy import histogram as histogram
 from keras.src.ops.numpy import hstack as hstack
 from keras.src.ops.numpy import identity as identity
diff --git a/keras/api/ops/numpy/__init__.py b/keras/api/ops/numpy/__init__.py
@@ -80,6 +80,7 @@
 from keras.src.ops.numpy import greater_equal as greater_equal
 from keras.src.ops.numpy import hamming as hamming
 from keras.src.ops.numpy import hanning as hanning
+from keras.src.ops.numpy import heaviside as heaviside
 from keras.src.ops.numpy import histogram as histogram
 from keras.src.ops.numpy import hstack as hstack
 from keras.src.ops.numpy import identity as identity
diff --git a/keras/src/backend/jax/numpy.py b/keras/src/backend/jax/numpy.py
@@ -52,6 +52,12 @@ def hanning(x):
     return jnp.hanning(x)
 
 
+def heaviside(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    return jnp.heaviside(x1, x2)
+
+
 def kaiser(x, beta):
     x = convert_to_tensor(x)
     return jnp.kaiser(x, beta)
diff --git a/keras/src/backend/numpy/numpy.py b/keras/src/backend/numpy/numpy.py
@@ -320,6 +320,19 @@ def hanning(x):
     return np.hanning(x).astype(config.floatx())
 
 
+def heaviside(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    if dtype in ["int8", "int16", "int32", "uint8", "uint16", "uint32"]:
+        dtype = config.floatx()
+    elif dtype in ["int64"]:
+        dtype = "float64"
+
+    return np.heaviside(x1, x2).astype(dtype)
+
+
 def kaiser(x, beta):
     x = convert_to_tensor(x)
     return np.kaiser(x, beta).astype(config.floatx())
diff --git a/keras/src/backend/openvino/excluded_concrete_tests.txt b/keras/src/backend/openvino/excluded_concrete_tests.txt
@@ -12,6 +12,7 @@ NumpyDtypeTest::test_bartlett
 NumpyDtypeTest::test_blackman
 NumpyDtypeTest::test_hamming
 NumpyDtypeTest::test_hanning
+NumpyDtypeTest::test_heaviside
 NumpyDtypeTest::test_kaiser
 NumpyDtypeTest::test_bitwise
 NumpyDtypeTest::test_cbrt
@@ -145,6 +146,7 @@ NumpyTwoInputOpsCorrectnessTest::test_cross
 NumpyTwoInputOpsCorrectnessTest::test_digitize
 NumpyTwoInputOpsCorrectnessTest::test_divide_no_nan
 NumpyTwoInputOpsCorrectnessTest::test_einsum
+NumpyTwoInputOpsCorrectnessTest::test_heaviside
 NumpyTwoInputOpsCorrectnessTest::test_inner
 NumpyTwoInputOpsCorrectnessTest::test_linspace
 NumpyTwoInputOpsCorrectnessTest::test_logspace
@@ -164,6 +166,8 @@ NumpyOneInputOpsDynamicShapeTest::test_kaiser
 NumpyOneInputOpsStaticShapeTest::test_angle
 NumpyOneInputOpsStaticShapeTest::test_cbrt
 NumpyOneInputOpsStaticShapeTest::test_deg2rad
+NumpyTwoInputOpsDynamicShapeTest::test_heaviside
+NumpyTwoInputOpsStaticShapeTest::test_heaviside
 CoreOpsBehaviorTests::test_associative_scan_invalid_arguments
 CoreOpsBehaviorTests::test_scan_invalid_arguments
 CoreOpsCallsTests::test_associative_scan_basic_call
diff --git a/keras/src/backend/openvino/numpy.py b/keras/src/backend/openvino/numpy.py
@@ -481,6 +481,12 @@ def hamming(x):
     )
 
 
+def heaviside(x1, x2):
+    raise NotImplementedError(
+        "`heaviside` is not supported with openvino backend"
+    )
+
+
 def kaiser(x, beta):
     raise NotImplementedError("`kaiser` is not supported with openvino backend")
 
diff --git a/keras/src/backend/tensorflow/numpy.py b/keras/src/backend/tensorflow/numpy.py
@@ -156,6 +156,26 @@ def hanning(x):
     return tf.signal.hann_window(x, periodic=False)
 
 
+def heaviside(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    if dtype in ["int8", "int16", "int32", "uint8", "uint16", "uint32"]:
+        dtype = config.floatx()
+    elif dtype in ["int64"]:
+        dtype = "float64"
+
+    x1 = tf.cast(x1, dtype)
+    x2 = tf.cast(x2, dtype)
+
+    return tf.where(
+        x1 < 0,
+        tf.zeros_like(x1),
+        tf.where(x1 > 0, tf.ones_like(x1), x2),
+    )
+
+
 def kaiser(x, beta):
     x = convert_to_tensor(x, dtype=tf.int32)
     return tf.signal.kaiser_window(x, beta=beta)
diff --git a/keras/src/backend/torch/numpy.py b/keras/src/backend/torch/numpy.py
@@ -445,6 +445,22 @@ def hanning(x):
     return torch.signal.windows.hann(x)
 
 
+def heaviside(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    if dtype in ["int8", "int16", "int32", "uint8", "uint16", "uint32"]:
+        dtype = config.floatx()
+    elif dtype == "int64":
+        dtype = "float64"
+
+    x1 = cast(x1, dtype)
+    x2 = cast(x2, dtype)
+
+    return torch.heaviside(x1, x2)
+
+
 def kaiser(x, beta):
     x = convert_to_tensor(x)
     return torch.signal.windows.kaiser(x, beta=beta)
diff --git a/keras/src/ops/numpy.py b/keras/src/ops/numpy.py
@@ -1317,6 +1317,44 @@ def hanning(x):
     return backend.numpy.hanning(x)
 
 
+class Heaviside(Operation):
+    def call(self, x1, x2):
+        return backend.numpy.heaviside(x1, x2)
+
+    def compute_output_spec(self, x1, x2):
+        dtype = dtypes.result_type(x1.dtype, x2.dtype)
+        if dtype in ["int8", "int16", "int32", "uint8", "uint16", "uint32"]:
+            dtype = backend.floatx()
+        elif dtype == "int64":
+            dtype = "float64"
+        return KerasTensor(broadcast_shapes(x1.shape, x2.shape), dtype=dtype)
+
+
+@keras_export(["keras.ops.heaviside", "keras.ops.numpy.heaviside"])
+def heaviside(x1, x2):
+    """Heaviside step function.
+
+    The Heaviside step function is defined as:
+    `heaviside(x1, x2) = 0 if x1 < 0, 1 if x1 > 0, x2 if x1 == 0`
+
+    Args:
+        x1: A tensor input.
+        x2: A scalar or tensor, the value to return when `x1 == 0`.
+
+    Returns:
+        A tensor with a shape determined by broadcasting `x1` and `x2`.
+
+    Example:
+    >>> x1 = keras.ops.convert_to_tensor([-2.0, 0.0, 3.0])
+    >>> x2 = 0.5
+    >>> keras.ops.heaviside(x1, x2)
+    array([0. , 0.5, 1. ], dtype=float32)
+    """
+    if any_symbolic_tensors((x1, x2)):
+        return Heaviside().symbolic_call(x1, x2)
+    return backend.numpy.heaviside(x1, x2)
+
+
 class Kaiser(Operation):
     def __init__(self, beta, *, name=None):
         super().__init__(name=name)
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
@@ -90,6 +90,11 @@ def test_add(self):
         y = KerasTensor((2, None))
         self.assertEqual(knp.add(x, y).shape, (2, 3))
 
+    def test_heaviside(self):
+        x = KerasTensor((None, 3))
+        y = KerasTensor((None, 3))
+        self.assertEqual(knp.heaviside(x, y).shape, (None, 3))
+
     def test_subtract(self):
         x = KerasTensor((None, 3))
         y = KerasTensor((2, None))
@@ -497,6 +502,19 @@ def test_add(self):
             y = KerasTensor((2, 3, 4))
             knp.add(x, y)
 
+    def test_heaviside(self):
+        x = KerasTensor((2, 3))
+        y = KerasTensor((2, 3))
+        self.assertEqual(knp.heaviside(x, y).shape, (2, 3))
+
+        x = KerasTensor((2, 3))
+        y = KerasTensor((3,))
+        self.assertEqual(knp.heaviside(x, y).shape, (2, 3))
+
+        x = KerasTensor((2, 3))
+        y = KerasTensor((1, 3))
+        self.assertEqual(knp.heaviside(x, y).shape, (2, 3))
+
     def test_subtract(self):
         x = KerasTensor((2, 3))
         y = KerasTensor((2, 3))
@@ -2342,6 +2360,17 @@ def test_add(self):
         self.assertAllClose(knp.Add()(x, y), np.add(x, y))
         self.assertAllClose(knp.Add()(x, z), np.add(x, z))
 
+    def test_heaviside(self):
+        x = np.array([[1, 2, 3], [3, 2, 1]])
+        y = np.array([[4, 5, 6], [6, 5, 4]])
+        self.assertAllClose(knp.heaviside(x, y), np.heaviside(x, y))
+        self.assertAllClose(knp.Heaviside()(x, y), np.heaviside(x, y))
+
+        x = np.array([[1, 2, 3], [3, 2, 1]])
+        y = np.array(4)
+        self.assertAllClose(knp.heaviside(x, y), np.heaviside(x, y))
+        self.assertAllClose(knp.Heaviside()(x, y), np.heaviside(x, y))
+
     def test_subtract(self):
         x = np.array([[1, 2, 3], [3, 2, 1]])
         y = np.array([[4, 5, 6], [3, 2, 1]])
@@ -7327,6 +7356,27 @@ def test_greater_equal(self, dtypes):
             expected_dtype,
         )
 
+    @parameterized.named_parameters(
+        named_product(dtypes=itertools.combinations(ALL_DTYPES, 2))
+    )
+    def test_heaviside(self, dtypes):
+        import jax.numpy as jnp
+
+        dtype1, dtype2 = dtypes
+        x1 = knp.ones((1, 1), dtype=dtype1)
+        x2 = knp.ones((1, 1), dtype=dtype2)
+        x1_jax = jnp.ones((1, 1), dtype=dtype1)
+        x2_jax = jnp.ones((1, 1), dtype=dtype2)
+        expected_dtype = standardize_dtype(jnp.heaviside(x1_jax, x2_jax).dtype)
+
+        self.assertEqual(
+            standardize_dtype(knp.heaviside(x1, x2).dtype), expected_dtype
+        )
+        self.assertEqual(
+            standardize_dtype(knp.Heaviside().symbolic_call(x1, x2).dtype),
+            expected_dtype,
+        )
+
     @parameterized.named_parameters(
         named_product(dtypes=itertools.combinations(ALL_DTYPES, 2))
     )
