[OpenVINO backend] support categorical

Author: Mohamed-Ashraf273

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -188,3 +188,58 @@ CoreOpsDtypeTest::test_convert_to_tensor12
 CoreOpsDtypeTest::test_convert_to_tensor14
 CoreOpsDtypeTest::test_convert_to_tensor25
 CoreOpsDtypeTest::test_convert_to_tensor37
+RandomCorrectnessTest::test_beta0
+RandomCorrectnessTest::test_beta1
+RandomCorrectnessTest::test_beta2
+RandomCorrectnessTest::test_binomial0
+RandomCorrectnessTest::test_binomial1
+RandomCorrectnessTest::test_binomial2
+RandomCorrectnessTest::test_dropout
+RandomCorrectnessTest::test_dropout_noise_shape
+RandomCorrectnessTest::test_gamma0
+RandomCorrectnessTest::test_gamma1
+RandomCorrectnessTest::test_gamma2
+RandomCorrectnessTest::test_randint0
+RandomCorrectnessTest::test_randint1
+RandomCorrectnessTest::test_randint2
+RandomCorrectnessTest::test_randint3
+RandomCorrectnessTest::test_randint4
+RandomCorrectnessTest::test_shuffle
+RandomCorrectnessTest::test_truncated_normal0
+RandomCorrectnessTest::test_truncated_normal1
+RandomCorrectnessTest::test_truncated_normal2
+RandomCorrectnessTest::test_truncated_normal3
+RandomCorrectnessTest::test_truncated_normal4
+RandomCorrectnessTest::test_truncated_normal5
+RandomCorrectnessTest::test_uniform0
+RandomCorrectnessTest::test_uniform1
+RandomCorrectnessTest::test_uniform2
+RandomCorrectnessTest::test_uniform3
+RandomCorrectnessTest::test_uniform4
+RandomBehaviorTest::test_beta_tf_data_compatibility
+RandomDTypeTest::test_beta_bfloat16
+RandomDTypeTest::test_beta_float16
+RandomDTypeTest::test_beta_float32
+RandomDTypeTest::test_beta_float64
+RandomDTypeTest::test_binomial_bfloat16
+RandomDTypeTest::test_binomial_float16
+RandomDTypeTest::test_binomial_float32
+RandomDTypeTest::test_binomial_float64
+RandomDTypeTest::test_dropout_bfloat16
+RandomDTypeTest::test_dropout_float16
+RandomDTypeTest::test_dropout_float32
+RandomDTypeTest::test_dropout_float64
+RandomDTypeTest::test_gamma_bfloat16
+RandomDTypeTest::test_gamma_float16
+RandomDTypeTest::test_gamma_float32
+RandomDTypeTest::test_gamma_float64
+RandomDTypeTest::test_normal_bfloat16
+RandomDTypeTest::test_randint_int16
+RandomDTypeTest::test_randint_int32
+RandomDTypeTest::test_randint_int64
+RandomDTypeTest::test_randint_int8
+RandomDTypeTest::test_randint_uint16
+RandomDTypeTest::test_randint_uint32
+RandomDTypeTest::test_randint_uint8
+RandomDTypeTest::test_truncated_normal_bfloat16
+RandomDTypeTest::test_uniform_bfloat16

keras/src/backend/openvino/excluded_tests.txt

@@ -33,7 +33,7 @@ keras/src/ops/math_test.py
 keras/src/ops/nn_test.py
 keras/src/optimizers
 keras/src/quantizers
-keras/src/random
+keras/src/random/seed_generator_test.py
 keras/src/regularizers
 keras/src/saving
 keras/src/trainers

keras/src/backend/openvino/random.py

@@ -6,6 +6,7 @@
 from keras.src.backend.openvino.core import OPENVINO_DTYPES
 from keras.src.backend.openvino.core import OpenVINOKerasTensor
 from keras.src.backend.openvino.core import convert_to_numpy
+from keras.src.backend.openvino.core import get_ov_output
 from keras.src.random.seed_generator import SeedGenerator
 from keras.src.random.seed_generator import draw_seed
 from keras.src.random.seed_generator import make_default_seed
@@ -39,9 +40,63 @@ def uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
 
 
 def categorical(logits, num_samples, dtype="int64", seed=None):
-    raise NotImplementedError(
-        "`categorical` is not supported with openvino backend"
+    def get_shape_dims(x):
+        shape = ov_opset.shape_of(x, Type.i32)
+        rank_tensor = ov_opset.shape_of(shape, Type.i32)
+        rank_scalar = ov_opset.squeeze(
+            rank_tensor, ov_opset.constant([0], Type.i32)
+        )
+        indices = ov_opset.range(
+            ov_opset.constant(0, Type.i32),
+            rank_scalar,
+            ov_opset.constant(1, Type.i32),
+            output_type=Type.i32,
+        )
+        return ov_opset.gather(shape, indices, axis=0)
+
+    dtype = dtype or "int64"
+    ov_dtype = OPENVINO_DTYPES[dtype]
+    logits = get_ov_output(logits)
+    probs = ov_opset.softmax(logits, axis=-1)
+    cumsum_probs = ov_opset.cumsum(probs, ov_opset.constant(-1, dtype="int32"))
+    shape = get_shape_dims(logits)
+    rank_tensor = ov_opset.shape_of(shape, Type.i32)
+    rank = ov_opset.squeeze(rank_tensor, ov_opset.constant([0], dtype=Type.i32))
+    rank_minus_1 = ov_opset.subtract(rank, ov_opset.constant(1, dtype=Type.i32))
+    indices = ov_opset.range(
+        ov_opset.constant(0, dtype=Type.i32),
+        rank_minus_1,
+        ov_opset.constant(1, dtype=Type.i32),
+        output_type=Type.i32,
+    )
+    batch_shape = ov_opset.gather(shape, indices, axis=0)
+    final_shape = ov_opset.concat(
+        [batch_shape, ov_opset.constant([num_samples], dtype=Type.i32)], axis=0
+    )
+    seed_tensor = draw_seed(seed)
+    if isinstance(seed_tensor, OpenVINOKerasTensor):
+        seed1, seed2 = convert_to_numpy(seed_tensor)
+    else:
+        seed1, seed2 = seed_tensor.data
+    rand = ov_opset.random_uniform(
+        final_shape,
+        ov_opset.constant(0.0, dtype=probs.get_element_type()),
+        ov_opset.constant(1.0, dtype=probs.get_element_type()),
+        probs.get_element_type(),
+        seed1,
+        seed2,
+    )
+    rand = ov_opset.unsqueeze(rand, [-1])
+    cumsum_probs = ov_opset.unsqueeze(
+        cumsum_probs, ov_opset.constant([1], dtype=Type.i32)
+    )
+    greater = ov_opset.greater(rand, cumsum_probs)
+    samples = ov_opset.reduce_sum(
+        ov_opset.convert(greater, Type.i32),
+        ov_opset.constant([-1], dtype=Type.i32),
     )
+    samples = ov_opset.convert(samples, ov_dtype)
+    return OpenVINOKerasTensor(samples.output(0))
 
 
 def randint(shape, minval, maxval, dtype="int32", seed=None):