[JAX FE]: Support jax.lax.eq and jax.lax.ne operation for JAX (openvi…

…notoolkit#26719) ### Details: - support jax.lax.eq and jax.lax.ne operation - enable unit tests for new operations ### Tickets: - [None](openvinotoolkit#26571)
sshlyapn · Sep 22, 2024 · f00ac41 · f00ac41
1 parent 94749db
commit f00ac41
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/src/frontends/jax/src/op/binary_op.cpp b/src/frontends/jax/src/op/binary_op.cpp
@@ -3,8 +3,10 @@
 //
 
 #include "openvino/frontend/jax/node_context.hpp"
+#include "openvino/op/equal.hpp"
 #include "openvino/op/greater.hpp"
 #include "openvino/op/greater_eq.hpp"
+#include "openvino/op/not_equal.hpp"
 #include "utils.hpp"
 
 using namespace std;
@@ -25,8 +27,10 @@ OutputVector translate_binary_op(const NodeContext& context) {
     return {binary_op};
 }
 
+template OutputVector translate_binary_op<v1::Equal>(const NodeContext& context);
 template OutputVector translate_binary_op<v1::GreaterEqual>(const NodeContext& context);
 template OutputVector translate_binary_op<v1::Greater>(const NodeContext& context);
+template OutputVector translate_binary_op<v1::NotEqual>(const NodeContext& context);
 
 }  // namespace op
 }  // namespace jax

diff --git a/src/frontends/jax/src/op_table.cpp b/src/frontends/jax/src/op_table.cpp
@@ -6,12 +6,14 @@
 
 #include "openvino/op/add.hpp"
 #include "openvino/op/divide.hpp"
+#include "openvino/op/equal.hpp"
 #include "openvino/op/erf.hpp"
 #include "openvino/op/exp.hpp"
 #include "openvino/op/greater.hpp"
 #include "openvino/op/greater_eq.hpp"
 #include "openvino/op/maximum.hpp"
 #include "openvino/op/multiply.hpp"
+#include "openvino/op/not_equal.hpp"
 #include "openvino/op/reduce_max.hpp"
 #include "openvino/op/reduce_sum.hpp"
 #include "openvino/op/sqrt.hpp"
@@ -63,13 +65,15 @@ const std::map<std::string, CreatorFunction> get_supported_ops_jaxpr() {
             {"device_put", op::skip_node},
             {"div", op::translate_1to1_match_2_inputs<v1::Divide>},
             {"dot_general", op::translate_dot_general},
+            {"eq", op::translate_binary_op<v1::Equal>},
             {"erf", op::translate_1to1_match_1_input<v0::Erf>},
             {"exp", op::translate_1to1_match_1_input<v0::Exp>},
             {"ge", op::translate_binary_op<v1::GreaterEqual>},
             {"gt", op::translate_binary_op<v1::Greater>},
             {"integer_pow", op::translate_integer_pow},
             {"max", op::translate_1to1_match_2_inputs<v1::Maximum>},
             {"mul", op::translate_1to1_match_2_inputs<v1::Multiply>},
+            {"ne", op::translate_binary_op<v1::NotEqual>},
             {"reduce_max", op::translate_reduce_op<v1::ReduceMax>},
             {"reduce_sum", op::translate_reduce_op<v1::ReduceSum>},
             {"reduce_window_max", op::translate_reduce_window_max},

diff --git a/tests/layer_tests/jax_tests/test_binary_comparison.py b/tests/layer_tests/jax_tests/test_binary_comparison.py
@@ -27,8 +27,10 @@ def _prepare_input(self):
 
     def create_model(self, input_shapes, binary_op, input_type):
         reduce_map = {
+            'eq': lax.eq,
             'ge': lax.ge,
-            'gt': lax.gt
+            'gt': lax.gt,
+            'ne': lax.ne
         }
 
         self.input_shapes = input_shapes
@@ -42,7 +44,7 @@ def jax_binary(x, y):
     @pytest.mark.parametrize('input_shapes', [[[5], [1]], [[1], [5]], [[2, 2, 4], [1, 1, 4]],
                                               [[5, 10], [5, 10]], [[2, 4, 6], [1, 4, 6]],
                                               [[5, 8, 10, 128], [5, 1, 10, 128]]])
-    @pytest.mark.parametrize('binary_op', ['ge', 'gt'])
+    @pytest.mark.parametrize('binary_op', ['eq', 'ge', 'gt', 'ne'])
     @pytest.mark.parametrize('input_type', [np.int8, np.uint8, np.int16, np.uint16,
                                             np.int32, np.uint32, np.int64, np.uint64,
                                             np.float16, np.float32, np.float64])