Add assume_unique and invert arguments to keras.ops.isin (#21552)

* Add arguments for isin method

* Add test cases

* Update argument description

* Add test cases more

Author: shashaka

keras/src/backend/jax/numpy.py

@@ -769,10 +769,10 @@ def isfinite(x):
     return jnp.isfinite(x)
 
 
-def isin(x1, x2):
+def isin(x1, x2, assume_unique=False, invert=False):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
-    return jnp.isin(x1, x2)
+    return jnp.isin(x1, x2, assume_unique=assume_unique, invert=invert)
 
 
 @sparse.elementwise_unary(linear=False)

keras/src/backend/numpy/numpy.py

@@ -682,10 +682,10 @@ def isfinite(x):
     return np.isfinite(x)
 
 
-def isin(x1, x2):
+def isin(x1, x2, assume_unique=False, invert=False):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
-    return np.isin(x1, x2)
+    return np.isin(x1, x2, assume_unique=assume_unique, invert=invert)
 
 
 def isinf(x):

keras/src/backend/openvino/numpy.py

@@ -929,7 +929,7 @@ def isfinite(x):
     return OpenVINOKerasTensor(ov_opset.is_finite(x).output(0))
 
 
-def isin(x1, x2):
+def isin(x1, x2, assume_unique=False, invert=False):
     raise NotImplementedError("`isin` is not supported with openvino backend")
 
 

keras/src/backend/tensorflow/numpy.py

@@ -1592,7 +1592,7 @@ def isfinite(x):
     return tf.math.is_finite(x)
 
 
-def isin(x1, x2):
+def isin(x1, x2, assume_unique=False, invert=False):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
 
@@ -1605,12 +1605,18 @@ def isin(x1, x2):
     x1 = tf.reshape(x1, [-1])
     x2 = tf.reshape(x2, [-1])
 
+    if not assume_unique:
+        x2 = tf.unique(x2)[0]
+
     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)
 
+    if invert:
+        result_flat = tf.logical_not(result_flat)
+
     return tf.reshape(result_flat, output_shape)
 
 

keras/src/backend/torch/numpy.py

@@ -886,7 +886,7 @@ def isfinite(x):
     return torch.isfinite(x)
 
 
-def isin(x1, x2):
+def isin(x1, x2, assume_unique=False, invert=False):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
 
@@ -900,7 +900,7 @@ def isin(x1, x2):
     if standardize_dtype(x2.dtype) == "bool":
         x2 = cast(x2, x1.dtype)
 
-    return torch.isin(x1, x2)
+    return torch.isin(x1, x2, assume_unique=assume_unique, invert=invert)
 
 
 def isinf(x):

keras/src/ops/numpy.py

@@ -3617,15 +3617,28 @@ def isfinite(x):
 
 
 class IsIn(Operation):
+    def __init__(
+        self,
+        assume_unique=False,
+        invert=False,
+        *,
+        name=None,
+    ):
+        super().__init__(name=name)
+        self.assume_unique = assume_unique
+        self.invert = invert
+
     def call(self, x1, x2):
-        return backend.numpy.isin(x1, x2)
+        return backend.numpy.isin(
+            x1, x2, assume_unique=self.assume_unique, invert=self.invert
+        )
 
     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):
+def isin(x1, x2, assume_unique=False, invert=False):
     """Test whether each element of `x1` is present in `x2`.
 
     This operation performs element-wise checks to determine if each value
@@ -3637,6 +3650,13 @@ def isin(x1, x2):
         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.
+        assume_unique: Boolean (default: False).
+            If True, assumes both `x1` and `x2` contain only unique elements.
+            This can speed up the computation. If False, duplicates will be
+            handled correctly but may impact performance.
+        invert: A boolean (default: False).
+            If True, inverts the result. Entries will be `True`
+            where `x1` elements are not in `x2`.
 
     Returns:
         A boolean tensor of the same shape as `x1` indicating element-wise
@@ -3650,8 +3670,12 @@ def 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)
+        return IsIn(assume_unique=assume_unique, invert=invert).symbolic_call(
+            x1, x2
+        )
+    return backend.numpy.isin(
+        x1, x2, assume_unique=assume_unique, invert=invert
+    )
 
 
 class Isinf(Operation):

keras/src/ops/numpy_test.py

@@ -2880,10 +2880,59 @@ def test_isin(self):
         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, assume_unique=True),
+            np.isin(x, y, assume_unique=True),
+        )
+        self.assertAllClose(
+            knp.isin(x, 2, assume_unique=True),
+            np.isin(x, 2, assume_unique=True),
+        )
+        self.assertAllClose(
+            knp.isin(2, x, assume_unique=True),
+            np.isin(2, x, assume_unique=True),
+        )
+
+        self.assertAllClose(
+            knp.isin(x, y, invert=True), np.isin(x, y, invert=True)
+        )
+        self.assertAllClose(
+            knp.isin(x, 2, invert=True), np.isin(x, 2, invert=True)
+        )
+        self.assertAllClose(
+            knp.isin(2, x, invert=True), np.isin(2, x, invert=True)
+        )
+
+        self.assertAllClose(
+            knp.isin(x, y, assume_unique=True, invert=True),
+            np.isin(x, y, assume_unique=True, invert=True),
+        )
+        self.assertAllClose(
+            knp.isin(x, 2, assume_unique=True, invert=True),
+            np.isin(x, 2, assume_unique=True, invert=True),
+        )
+        self.assertAllClose(
+            knp.isin(2, x, assume_unique=True, invert=True),
+            np.isin(2, x, assume_unique=True, invert=True),
+        )
+
         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(assume_unique=True)(x, y),
+            np.isin(x, y, assume_unique=True),
+        )
+        self.assertAllClose(
+            knp.IsIn(invert=True)(x, y),
+            np.isin(x, y, invert=True),
+        )
+        self.assertAllClose(
+            knp.IsIn(assume_unique=True, invert=True)(x, y),
+            np.isin(x, y, assume_unique=True, invert=True),
+        )
+
     def test_less(self):
         x = np.array([[1, 2, 3], [3, 2, 1]])
         y = np.array([[4, 5, 6], [3, 2, 1]])