Implement kaiser function in keras.ops (#21303)

* Add kaiser window for ops

* Add test cases

* update excluded_concrete_tests

Author: shashaka

keras/api/_tf_keras/keras/ops/__init__.py

@@ -195,6 +195,7 @@
 from keras.src.ops.numpy import isfinite as isfinite
 from keras.src.ops.numpy import isinf as isinf
 from keras.src.ops.numpy import isnan as isnan
+from keras.src.ops.numpy import kaiser as kaiser
 from keras.src.ops.numpy import left_shift as left_shift
 from keras.src.ops.numpy import less as less
 from keras.src.ops.numpy import less_equal as less_equal

keras/api/_tf_keras/keras/ops/numpy/__init__.py

@@ -85,6 +85,7 @@
 from keras.src.ops.numpy import isfinite as isfinite
 from keras.src.ops.numpy import isinf as isinf
 from keras.src.ops.numpy import isnan as isnan
+from keras.src.ops.numpy import kaiser as kaiser
 from keras.src.ops.numpy import left_shift as left_shift
 from keras.src.ops.numpy import less as less
 from keras.src.ops.numpy import less_equal as less_equal

keras/api/ops/__init__.py

@@ -195,6 +195,7 @@
 from keras.src.ops.numpy import isfinite as isfinite
 from keras.src.ops.numpy import isinf as isinf
 from keras.src.ops.numpy import isnan as isnan
+from keras.src.ops.numpy import kaiser as kaiser
 from keras.src.ops.numpy import left_shift as left_shift
 from keras.src.ops.numpy import less as less
 from keras.src.ops.numpy import less_equal as less_equal

keras/api/ops/numpy/__init__.py

@@ -85,6 +85,7 @@
 from keras.src.ops.numpy import isfinite as isfinite
 from keras.src.ops.numpy import isinf as isinf
 from keras.src.ops.numpy import isnan as isnan
+from keras.src.ops.numpy import kaiser as kaiser
 from keras.src.ops.numpy import left_shift as left_shift
 from keras.src.ops.numpy import less as less
 from keras.src.ops.numpy import less_equal as less_equal

keras/src/backend/jax/numpy.py

@@ -47,6 +47,11 @@ def hamming(x):
     return jnp.hamming(x)
 
 
+def kaiser(x, beta):
+    x = convert_to_tensor(x)
+    return jnp.kaiser(x, beta)
+
+
 def bincount(x, weights=None, minlength=0, sparse=False):
     # Note: bincount is never tracable / jittable because the output shape
     # depends on the values in x.

keras/src/backend/numpy/numpy.py

@@ -315,6 +315,11 @@ def hamming(x):
     return np.hamming(x).astype(config.floatx())
 
 
+def kaiser(x, beta):
+    x = convert_to_tensor(x)
+    return np.kaiser(x, beta).astype(config.floatx())
+
+
 def bincount(x, weights=None, minlength=0, sparse=False):
     if sparse:
         raise ValueError("Unsupported value `sparse=True` with numpy backend")

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -11,6 +11,7 @@ NumpyDtypeTest::test_array
 NumpyDtypeTest::test_bartlett
 NumpyDtypeTest::test_blackman
 NumpyDtypeTest::test_hamming
+NumpyDtypeTest::test_kaiser
 NumpyDtypeTest::test_bitwise
 NumpyDtypeTest::test_ceil
 NumpyDtypeTest::test_concatenate
@@ -75,6 +76,7 @@ NumpyOneInputOpsCorrectnessTest::test_array
 NumpyOneInputOpsCorrectnessTest::test_bartlett
 NumpyOneInputOpsCorrectnessTest::test_blackman
 NumpyOneInputOpsCorrectnessTest::test_hamming
+NumpyOneInputOpsCorrectnessTest::test_kaiser
 NumpyOneInputOpsCorrectnessTest::test_bitwise_invert
 NumpyOneInputOpsCorrectnessTest::test_conj
 NumpyOneInputOpsCorrectnessTest::test_correlate
@@ -148,4 +150,5 @@ NumpyOneInputOpsDynamicShapeTest::test_angle
 NumpyOneInputOpsDynamicShapeTest::test_bartlett
 NumpyOneInputOpsDynamicShapeTest::test_blackman
 NumpyOneInputOpsDynamicShapeTest::test_hamming
+NumpyOneInputOpsDynamicShapeTest::test_kaiser
 NumpyOneInputOpsStaticShapeTest::test_angle

keras/src/backend/openvino/numpy.py

@@ -481,6 +481,10 @@ def hamming(x):
     )
 
 
+def kaiser(x, beta):
+    raise NotImplementedError("`kaiser` is not supported with openvino backend")
+
+
 def bincount(x, weights=None, minlength=0, sparse=False):
     if x is None:
         raise ValueError("input x is None")

keras/src/backend/tensorflow/numpy.py

@@ -151,6 +151,11 @@ def hamming(x):
     return tf.signal.hamming_window(x, periodic=False)
 
 
+def kaiser(x, beta):
+    x = convert_to_tensor(x, dtype=tf.int32)
+    return tf.signal.kaiser_window(x, beta=beta)
+
+
 def bincount(x, weights=None, minlength=0, sparse=False):
     x = convert_to_tensor(x)
     dtypes_to_resolve = [x.dtype]

keras/src/backend/torch/numpy.py

@@ -440,6 +440,11 @@ def hamming(x):
     return torch.signal.windows.hamming(x)
 
 
+def kaiser(x, beta):
+    x = convert_to_tensor(x)
+    return torch.signal.windows.kaiser(x, beta=beta)
+
+
 def bincount(x, weights=None, minlength=0, sparse=False):
     if sparse:
         raise ValueError("Unsupported value `sparse=True` with torch backend")

keras/src/ops/numpy.py

@@ -1270,6 +1270,44 @@ def hamming(x):
     return backend.numpy.hamming(x)
 
 
+class Kaiser(Operation):
+    def __init__(self, beta):
+        super().__init__()
+        self.beta = beta
+
+    def call(self, x):
+        return backend.numpy.kaiser(x, self.beta)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=backend.floatx())
+
+
+@keras_export(["keras.ops.kaiser", "keras.ops.numpy.kaiser"])
+def kaiser(x, beta):
+    """Kaiser window function.
+
+    The Kaiser window is defined as:
+    `w[n] = I0(beta * sqrt(1 - (2n / (N - 1) - 1)^2)) / I0(beta)`
+    where I0 is the modified zeroth-order Bessel function of the first kind.
+
+    Args:
+        x: Scalar or 1D Tensor. The window length.
+        beta: Float. Shape parameter for the Kaiser window.
+
+    Returns:
+        A 1D tensor containing the Kaiser window values.
+
+    Example:
+    >>> x = keras.ops.convert_to_tensor(5)
+    >>> keras.ops.kaiser(x, beta=14.0)
+    array([7.7268669e-06, 1.6493219e-01, 1.0000000e+00, 1.6493219e-01,
+       7.7268669e-06], dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return Kaiser(beta).symbolic_call(x)
+    return backend.numpy.kaiser(x, beta)
+
+
 class Bincount(Operation):
     def __init__(self, weights=None, minlength=0, sparse=False):
         super().__init__()

keras/src/ops/numpy_test.py

@@ -1232,6 +1232,11 @@ def test_hamming(self):
         x = np.random.randint(1, 100 + 1)
         self.assertEqual(knp.hamming(x).shape[0], x)
 
+    def test_kaiser(self):
+        x = np.random.randint(1, 100 + 1)
+        beta = float(np.random.randint(10, 20 + 1))
+        self.assertEqual(knp.kaiser(x, beta).shape[0], x)
+
     def test_bitwise_invert(self):
         x = KerasTensor((None, 3))
         self.assertEqual(knp.bitwise_invert(x).shape, (None, 3))
@@ -3620,6 +3625,13 @@ def test_hamming(self):
 
         self.assertAllClose(knp.Hamming()(x), np.hamming(x))
 
+    def test_kaiser(self):
+        x = np.random.randint(1, 100 + 1)
+        beta = float(np.random.randint(10, 20 + 1))
+        self.assertAllClose(knp.kaiser(x, beta), np.kaiser(x, beta))
+
+        self.assertAllClose(knp.Kaiser(beta)(x), np.kaiser(x, beta))
+
     @parameterized.named_parameters(
         named_product(sparse_input=(False, True), sparse_arg=(False, True))
     )
@@ -5627,6 +5639,24 @@ def test_hamming(self, dtype):
             expected_dtype,
         )
 
+    @parameterized.named_parameters(named_product(dtype=ALL_DTYPES))
+    def test_kaiser(self, dtype):
+        import jax.numpy as jnp
+
+        x = knp.ones((), dtype=dtype)
+        beta = knp.ones((), dtype=dtype)
+        x_jax = jnp.ones((), dtype=dtype)
+        beta_jax = jnp.ones((), dtype=dtype)
+        expected_dtype = standardize_dtype(jnp.kaiser(x_jax, beta_jax).dtype)
+
+        self.assertEqual(
+            standardize_dtype(knp.kaiser(x, beta).dtype), expected_dtype
+        )
+        self.assertEqual(
+            standardize_dtype(knp.Kaiser(beta).symbolic_call(x).dtype),
+            expected_dtype,
+        )
+
     @parameterized.named_parameters(named_product(dtype=INT_DTYPES))
     def test_bincount(self, dtype):
         import jax.numpy as jnp