Implement isin function in keras.ops (#21523)

shashaka · web-flow · commit 92aefef06640 · 2025-07-30T17:00:10.000-07:00
* Add isin method for numpy

* Add isin method for openvino

* Correct logic for tensorflow

* Add isin method for ops

* Add test cases and update description

* update excluded_concrete_tests.txt

* correct isin method for tensorflow
diff --git a/keras/api/_tf_keras/keras/ops/__init__.py b/keras/api/_tf_keras/keras/ops/__init__.py
@@ -199,6 +199,7 @@
 from keras.src.ops.numpy import inner as inner
 from keras.src.ops.numpy import isclose as isclose
 from keras.src.ops.numpy import isfinite as isfinite
+from keras.src.ops.numpy import isin as isin
 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
diff --git a/keras/api/_tf_keras/keras/ops/numpy/__init__.py b/keras/api/_tf_keras/keras/ops/numpy/__init__.py
@@ -88,6 +88,7 @@
 from keras.src.ops.numpy import inner as inner
 from keras.src.ops.numpy import isclose as isclose
 from keras.src.ops.numpy import isfinite as isfinite
+from keras.src.ops.numpy import isin as isin
 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
diff --git a/keras/api/ops/__init__.py b/keras/api/ops/__init__.py
@@ -199,6 +199,7 @@
 from keras.src.ops.numpy import inner as inner
 from keras.src.ops.numpy import isclose as isclose
 from keras.src.ops.numpy import isfinite as isfinite
+from keras.src.ops.numpy import isin as isin
 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
diff --git a/keras/api/ops/numpy/__init__.py b/keras/api/ops/numpy/__init__.py
@@ -88,6 +88,7 @@
 from keras.src.ops.numpy import inner as inner
 from keras.src.ops.numpy import isclose as isclose
 from keras.src.ops.numpy import isfinite as isfinite
+from keras.src.ops.numpy import isin as isin
 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
diff --git a/keras/src/backend/jax/numpy.py b/keras/src/backend/jax/numpy.py
@@ -769,6 +769,12 @@ def isfinite(x):
     return jnp.isfinite(x)
 
 
+def isin(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    return jnp.isin(x1, x2)
+
+
 @sparse.elementwise_unary(linear=False)
 def isinf(x):
     x = convert_to_tensor(x)
diff --git a/keras/src/backend/numpy/numpy.py b/keras/src/backend/numpy/numpy.py
@@ -682,6 +682,12 @@ def isfinite(x):
     return np.isfinite(x)
 
 
+def isin(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    return np.isin(x1, x2)
+
+
 def isinf(x):
     return np.isinf(x)
 
diff --git a/keras/src/backend/openvino/excluded_concrete_tests.txt b/keras/src/backend/openvino/excluded_concrete_tests.txt
@@ -32,6 +32,7 @@ NumpyDtypeTest::test_flip
 NumpyDtypeTest::test_floor
 NumpyDtypeTest::test_inner
 NumpyDtypeTest::test_isfinite
+NumpyDtypeTest::test_isin
 NumpyDtypeTest::test_isinf
 NumpyDtypeTest::test_isnan
 NumpyDtypeTest::test_linspace
@@ -142,6 +143,7 @@ NumpyTwoInputOpsCorrectnessTest::test_divide_no_nan
 NumpyTwoInputOpsCorrectnessTest::test_einsum
 NumpyTwoInputOpsCorrectnessTest::test_heaviside
 NumpyTwoInputOpsCorrectnessTest::test_inner
+NumpyTwoInputOpsCorrectnessTest::test_isin
 NumpyTwoInputOpsCorrectnessTest::test_linspace
 NumpyTwoInputOpsCorrectnessTest::test_logspace
 NumpyTwoInputOpsCorrectnessTest::test_quantile
@@ -160,7 +162,9 @@ NumpyOneInputOpsStaticShapeTest::test_angle
 NumpyOneInputOpsStaticShapeTest::test_cbrt
 NumpyOneInputOpsStaticShapeTest::test_deg2rad
 NumpyTwoInputOpsDynamicShapeTest::test_heaviside
+NumpyTwoInputOpsDynamicShapeTest::test_isin
 NumpyTwoInputOpsStaticShapeTest::test_heaviside
+NumpyTwoInputOpsStaticShapeTest::test_isin
 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
@@ -929,6 +929,10 @@ def isfinite(x):
     return OpenVINOKerasTensor(ov_opset.is_finite(x).output(0))
 
 
+def isin(x1, x2):
+    raise NotImplementedError("`isin` is not supported with openvino backend")
+
+
 def isinf(x):
     x = get_ov_output(x)
     return OpenVINOKerasTensor(ov_opset.is_inf(x).output(0))
diff --git a/keras/src/backend/tensorflow/numpy.py b/keras/src/backend/tensorflow/numpy.py
@@ -1592,6 +1592,28 @@ def isfinite(x):
     return tf.math.is_finite(x)
 
 
+def isin(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    x1 = tf.cast(x1, dtype)
+    x2 = tf.cast(x2, dtype)
+
+    output_shape = tf.shape(x1)
+
+    x1 = tf.reshape(x1, [-1])
+    x2 = tf.reshape(x2, [-1])
+
+    if tf.size(x1) == 0 or tf.size(x2) == 0:
+        return tf.zeros(output_shape, dtype=tf.bool)
+
+    cmp = tf.equal(tf.expand_dims(x1, 1), tf.expand_dims(x2, 0))
+    result_flat = tf.reduce_any(cmp, axis=1)
+
+    return tf.reshape(result_flat, output_shape)
+
+
 def isinf(x):
     x = convert_to_tensor(x)
     dtype_as_dtype = tf.as_dtype(x.dtype)
diff --git a/keras/src/backend/torch/numpy.py b/keras/src/backend/torch/numpy.py
@@ -886,6 +886,23 @@ def isfinite(x):
     return torch.isfinite(x)
 
 
+def isin(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    if dtype == "bool":
+        x1 = cast(x1, "int32")
+        x2 = cast(x2, "int32")
+
+    if standardize_dtype(x1.dtype) == "bool":
+        x1 = cast(x1, x2.dtype)
+    if standardize_dtype(x2.dtype) == "bool":
+        x2 = cast(x2, x1.dtype)
+
+    return torch.isin(x1, x2)
+
+
 def isinf(x):
     x = convert_to_tensor(x)
     return torch.isinf(x)
diff --git a/keras/src/ops/numpy.py b/keras/src/ops/numpy.py
@@ -3616,6 +3616,44 @@ def isfinite(x):
     return backend.numpy.isfinite(x)
 
 
+class IsIn(Operation):
+    def call(self, x1, x2):
+        return backend.numpy.isin(x1, x2)
+
+    def compute_output_spec(self, x1, x2):
+        return KerasTensor(x1.shape, dtype="bool")
+
+
+@keras_export(["keras.ops.isin", "keras.ops.numpy.isin"])
+def isin(x1, x2):
+    """Test whether each element of `x1` is present in `x2`.
+
+    This operation performs element-wise checks to determine if each value
+    in `x1` is contained within `x2`. The result is a boolean tensor with
+    the same shape as `x1`, where each entry is `True` if the corresponding
+    element in `x1` is in `x2`, and `False` otherwise.
+
+    Args:
+        x1: Input tensor or array-like structure to test.
+        x2: Values against which each element of `x1` is tested.
+            Can be a tensor, list, or scalar.
+
+    Returns:
+        A boolean tensor of the same shape as `x1` indicating element-wise
+        membership in `x2`.
+
+    Example:
+    >>> from keras import ops
+    >>> x1 = ops.array([0, 1, 2, 5])
+    >>> x2 = ops.array([0, 2])
+    >>> result = ops.isin(x1, x2)
+    array([ True, False,  True, False])
+    """
+    if any_symbolic_tensors((x1, x2)):
+        return IsIn().symbolic_call(x1, x2)
+    return backend.numpy.isin(x1, x2)
+
+
 class Isinf(Operation):
     def call(self, x):
         return backend.numpy.isinf(x)
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
@@ -230,6 +230,11 @@ def test_isclose(self):
         y = KerasTensor((2, None))
         self.assertEqual(knp.isclose(x, y).shape, (2, 3))
 
+    def test_isin(self):
+        x = KerasTensor((None, 3))
+        y = KerasTensor((2, None))
+        self.assertEqual(knp.isin(x, y).shape, (None, 3))
+
     def test_less(self):
         x = KerasTensor((None, 3))
         y = KerasTensor((2, None))
@@ -769,6 +774,14 @@ def test_isclose(self):
             y = KerasTensor((2, 3, 4))
             knp.isclose(x, y)
 
+    def test_isin(self):
+        x = KerasTensor((2, 3))
+        y = KerasTensor((3, 3))
+        self.assertEqual(knp.isin(x, y).shape, (2, 3))
+
+        x = KerasTensor((2, 3))
+        self.assertEqual(knp.isin(x, 2).shape, (2, 3))
+
     def test_less(self):
         x = KerasTensor((2, 3))
         y = KerasTensor((2, 3))
@@ -2852,6 +2865,17 @@ def test_isclose(self):
         self.assertAllClose(knp.Isclose()(x, 2), np.isclose(x, 2))
         self.assertAllClose(knp.Isclose()(2, x), np.isclose(2, x))
 
+    def test_isin(self):
+        x = np.array([[1, 2, 3], [3, 2, 1]])
+        y = np.array([[4, 5, 6], [3, 2, 1]])
+        self.assertAllClose(knp.isin(x, y), np.isin(x, y))
+        self.assertAllClose(knp.isin(x, 2), np.isin(x, 2))
+        self.assertAllClose(knp.isin(2, x), np.isin(2, x))
+
+        self.assertAllClose(knp.IsIn()(x, y), np.isin(x, y))
+        self.assertAllClose(knp.IsIn()(x, 2), np.isin(x, 2))
+        self.assertAllClose(knp.IsIn()(2, x), np.isin(2, x))
+
     def test_less(self):
         x = np.array([[1, 2, 3], [3, 2, 1]])
         y = np.array([[4, 5, 6], [3, 2, 1]])
@@ -7446,6 +7470,27 @@ def test_isfinite(self, dtype):
             expected_dtype,
         )
 
+    @parameterized.named_parameters(
+        named_product(dtypes=itertools.combinations(ALL_DTYPES, 2))
+    )
+    def test_isin(self, dtypes):
+        import jax.numpy as jnp
+
+        dtype1, dtype2 = dtypes
+        x1 = knp.ones((), dtype=dtype1)
+        x2 = knp.ones((), dtype=dtype2)
+        x1_jax = jnp.ones((), dtype=dtype1)
+        x2_jax = jnp.ones((), dtype=dtype2)
+        expected_dtype = standardize_dtype(jnp.isin(x1_jax, x2_jax).dtype)
+
+        self.assertEqual(
+            standardize_dtype(knp.isin(x1, x2).dtype), expected_dtype
+        )
+        self.assertEqual(
+            standardize_dtype(knp.IsIn().symbolic_call(x1, x2).dtype),
+            expected_dtype,
+        )
+
     @parameterized.named_parameters(named_product(dtype=ALL_DTYPES))
     def test_isinf(self, dtype):
         import jax.numpy as jnp
