[JAX FE]: add support for jax.lax.logistic

src/frontends/jax/src/op_table.cpp

@@ -20,6 +20,7 @@
 #include "openvino/op/reduce_max.hpp"
 #include "openvino/op/reduce_sum.hpp"
 #include "openvino/op/sqrt.hpp"
+#include "openvino/op/sigmoid.hpp"
 #include "openvino/op/subtract.hpp"
 #include "openvino/op/tanh.hpp"
 #include "utils.hpp"
@@ -53,6 +54,7 @@ OP_CONVERTER(translate_reduce_window_sum);
 OP_CONVERTER(translate_reshape);
 OP_CONVERTER(translate_rsqrt);
 OP_CONVERTER(translate_select_n);
+OP_CONVERTER(translate_logistic);
 OP_CONVERTER(translate_slice);
 OP_CONVERTER(translate_square);
 OP_CONVERTER(translate_squeeze);
@@ -94,6 +96,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_jaxpr() {
             {"rsqrt", op::translate_rsqrt},
             {"reshape", op::translate_reshape},
             {"select_n", op::translate_select_n},
+            {"logistic", op::translate_1to1_match_1_input<v0::Sigmoid>},
             {"slice", op::translate_slice},
             {"square", op::translate_square},
             {"sqrt", op::translate_1to1_match_1_input<v0::Sqrt>},

tests/layer_tests/jax_tests/test_logistic.py

@@ -0,0 +1,37 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import jax
+import numpy as np
+import pytest
+from jax import numpy as jnp
+
+from jax_layer_test_class import JaxLayerTest
+
+rng = np.random.default_rng(5402)
+
+
+class TestLogistic(JaxLayerTest):
+    def _prepare_input(self):
+
+        input = jnp.array(np.random.uniform(-1000, 1000, self.input_shape).astype(self.input_type))
+        return [input]
+
+    def create_model(self, input_shape, input_type):
+        self.input_shape = input_shape
+        self.input_type = input_type
+
+        def jax_logistic(input):
+            return jax.lax.logistic(input)
+
+        return jax_logistic, None, 'logistic'
+
+    @pytest.mark.parametrize("input_shape", [[2], [3, 4], [5,6,7]])
+    @pytest.mark.parametrize("input_type", [np.float32, np.float64])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_jax_fe
+    def test_logistic(self, ie_device, precision, ir_version, input_shape, input_type):
+        self._test(*self.create_model(input_shape, input_type),
+                   ie_device, precision,
+                   ir_version)