[TF FE][JAX FE] Support latest TF 2.18, JAX 0.4.35 and NumPy 2.x (openvinotoolkit#27246)

**Details:** Support TF 2.18 and JAX 0.4.35

**Ticket:** TBD

---------

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>

Author: rkazants

.github/workflows/job_tokenizers.yml

@@ -20,12 +20,15 @@ on:
         description: 'Components that are affected by changes in the commit defined by the Smart CI Action'
         type: string
         required: true
+      python-version:
+        description: 'Python version to setup. E.g., "3.11"'
+        type: string
+        required: true
 
 permissions: read-all
 
 env:
   PIP_CACHE_PATH: /mount/caches/pip/linux
-  PYTHON_VERSION: '3.11'
   TARGET_BRANCH: ${{ github.base_ref || github.event.merge_group.base_ref || github.ref }}
 
 jobs:
@@ -63,10 +66,10 @@ jobs:
         if: runner.os == 'macOS'
         run: brew install pigz
 
-      - name: Setup Python ${{ env.PYTHON_VERSION }}
+      - name: Setup Python ${{ inputs.python-version }}
         uses: ./.github/actions/setup_python
         with:
-          version: ${{ env.PYTHON_VERSION }}
+          version: ${{ inputs.python-version }}
           pip-cache-path: ${{ runner.os == 'Linux' && env.PIP_CACHE_PATH || '' }}
           should-setup-pip-paths: ${{ runner.os == 'Linux' }}
           self-hosted-runner: ${{ runner.os == 'Linux' }}

.github/workflows/linux_arm64.yml

@@ -146,6 +146,7 @@ jobs:
       shell: bash
       container: '{"image": "${{ fromJSON(needs.docker.outputs.images).ov_build.ubuntu_20_04_arm64 }}", "volumes": ["/mount:/mount"]}'
       affected-components: ${{ needs.smart_ci.outputs.affected_components }}
+      python-version: '3.11'
     if: fromJSON(needs.smart_ci.outputs.affected_components).TOKENIZERS
 
   CXX_Unit_Tests:

.github/workflows/mac.yml

@@ -258,6 +258,7 @@ jobs:
       runner: 'macos-13'
       shell: bash
       affected-components: ${{ needs.smart_ci.outputs.affected_components }}
+      python-version: '3.11'
     if: fromJSON(needs.smart_ci.outputs.affected_components).TOKENIZERS
 
   CXX_Unit_Tests:

.github/workflows/mac_arm64.yml

@@ -258,6 +258,7 @@ jobs:
       runner: 'macos-13-xlarge'
       shell: bash
       affected-components: ${{ needs.smart_ci.outputs.affected_components }}
+      python-version: '3.11'
     if: fromJSON(needs.smart_ci.outputs.affected_components).TOKENIZERS
 
   CXX_Unit_Tests:

.github/workflows/ubuntu_22.yml

@@ -491,6 +491,7 @@ jobs:
       shell: bash
       container: '{"image": "${{ fromJSON(needs.docker.outputs.images).ov_build.ubuntu_22_04_x64 }}", "volumes": ["/mount:/mount"]}'
       affected-components: ${{ needs.smart_ci.outputs.affected_components }}
+      python-version: '3.11'
     if: fromJSON(needs.smart_ci.outputs.affected_components).TOKENIZERS
 
   iGPU:

.github/workflows/ubuntu_24.yml

@@ -144,9 +144,31 @@ jobs:
       affected-components: ${{ needs.smart_ci.outputs.affected_components }}
       python-version: '3.12'
 
+  TensorFlow_Layer_Tests:
+    name: TensorFlow Layer Tests
+    needs: [ Docker, Build, Smart_CI, Openvino_tokenizers ]
+    uses: ./.github/workflows/job_tensorflow_layer_tests.yml
+    with:
+      runner: 'aks-linux-4-cores-16gb'
+      container: '{"image": "${{ fromJSON(needs.docker.outputs.images).ov_test.ubuntu_24_04_x64 }}", "volumes": ["/mount:/mount"]}'
+      affected-components: ${{ needs.smart_ci.outputs.affected_components }}
+      python-version: '3.12'
+
+  Openvino_tokenizers:
+    name: OpenVINO tokenizers extension
+    needs: [ Build, Smart_CI, Docker ]
+    uses: ./.github/workflows/job_tokenizers.yml
+    with:
+      runner: 'aks-linux-4-cores-16gb'
+      shell: bash
+      container: '{"image": "${{ fromJSON(needs.docker.outputs.images).ov_build.ubuntu_24_04_x64 }}", "volumes": ["/mount:/mount"]}'
+      affected-components: ${{ needs.smart_ci.outputs.affected_components }}
+      python-version: '3.12'
+    if: fromJSON(needs.smart_ci.outputs.affected_components).TOKENIZERS
+
   Overall_Status:
     name: ci/gha_overall_status_ubuntu_24
-    needs: [Smart_CI, Build, Debian_Packages, Samples, Python_Unit_Tests]
+    needs: [Smart_CI, Build, Debian_Packages, Samples, Python_Unit_Tests, Pytorch_Layer_Tests, TensorFlow_Layer_Tests, Openvino_tokenizers]
     if: ${{ always() }}
     runs-on: ubuntu-latest
     steps:

.github/workflows/windows_vs2019_release.yml

@@ -235,6 +235,7 @@ jobs:
       runner: 'aks-win-4-cores-8gb'
       shell: pwsh
       affected-components: ${{ needs.smart_ci.outputs.affected_components }}
+      python-version: '3.11'
     if: fromJSON(needs.smart_ci.outputs.affected_components).TOKENIZERS
 
   Python_Unit_Tests:

src/frontends/jax/src/op/erfc.cpp

@@ -0,0 +1,33 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/jax/node_context.hpp"
+#include "openvino/op/erf.hpp"
+#include "openvino/op/subtract.hpp"
+#include "utils.hpp"
+
+using namespace std;
+using namespace ov;
+using namespace ov::op;
+
+namespace ov {
+namespace frontend {
+namespace jax {
+namespace op {
+
+OutputVector translate_erfc(const NodeContext& context) {
+    num_inputs_check(context, 1, 1);
+    auto x = context.get_input(0);
+
+    // create const one of the same type as x
+    auto const_one = create_same_type_const_scalar<int64_t>(x, 1);
+    Output<Node> res = make_shared<v0::Erf>(x);
+    res = make_shared<v1::Subtract>(const_one, res);
+    return {res};
+};
+
+}  // namespace op
+}  // namespace jax
+}  // namespace frontend
+}  // namespace ov

src/frontends/jax/src/op_table.cpp

@@ -45,6 +45,7 @@ OP_CONVERTER(translate_convert);
 OP_CONVERTER(translate_convolution);
 OP_CONVERTER(translate_copy);
 OP_CONVERTER(translate_dot_general);
+OP_CONVERTER(translate_erfc);
 OP_CONVERTER(translate_integer_pow);
 OP_T_CONVERTER(translate_reduce_op);
 OP_CONVERTER(translate_reduce_window_max);
@@ -72,6 +73,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_jaxpr() {
             {"dot_general", op::translate_dot_general},
             {"eq", op::translate_binary_op<v1::Equal>},
             {"erf", op::translate_1to1_match_1_input<v0::Erf>},
+            {"erfc", op::translate_erfc},
             {"exp", op::translate_1to1_match_1_input<v0::Exp>},
             {"ge", op::translate_binary_op<v1::GreaterEqual>},
             {"gt", op::translate_binary_op<v1::Greater>},

tests/constraints.txt

@@ -11,7 +11,7 @@ sympy>=1.10
 wheel>=0.38.1
 defusedxml>=0.7.1
 fastjsonschema~=2.17.1
-tensorflow>=2.5,<2.18.0
+tensorflow>=2.5,<2.19.0
 requests>=2.25.1
 opencv-python>=4.5
 paddlepaddle==2.6.1
@@ -21,11 +21,11 @@ pytest>=5.0,<8.4
 pytest-dependency==0.5.1
 pytest-html==4.1.1
 pytest-timeout==2.3.1
-jax<=0.4.33
-jaxlib<=0.4.33
+jax<=0.4.35
+jaxlib<=0.4.35
 kornia==0.7.0
 networkx<=3.3
-flax<=0.9.0
+flax<=0.10.0
 
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch~=2.4.1; platform_system != "Darwin" or platform_machine != "x86_64"

tests/layer_tests/jax_tests/test_erfc.py

@@ -0,0 +1,38 @@
+# 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(109734)
+
+
+class TestErfc(JaxLayerTest):
+    def _prepare_input(self):
+        # erf are mostly changing in a range [-4, 4]
+        x = rng.uniform(-4.0, 4.0, self.input_shape).astype(self.input_type)
+
+        x = jnp.array(x)
+        return [x]
+
+    def create_model(self, input_shape, input_type):
+        self.input_shape = input_shape
+        self.input_type = input_type
+
+        def jax_erfc(x):
+            return jax.lax.erfc(x)
+
+        return jax_erfc, None, 'erfc'
+
+    @pytest.mark.parametrize("input_shape", [[2], [3, 4]])
+    @pytest.mark.parametrize("input_type", [np.float16, np.float32, np.float64])
+    @pytest.mark.nightly
+    @pytest.mark.precommit_jax_fe
+    def test_erfc(self, ie_device, precision, ir_version, input_shape, input_type):
+        self._test(*self.create_model(input_shape, input_type),
+                   ie_device, precision,
+                   ir_version)

tests/layer_tests/tensorflow_tests/test_tf_Equal.py

@@ -146,8 +146,8 @@ def test_tf_equal_int64(self, params, ie_device, precision, ir_version, temp_dir
 
     # Values for checking important corner cases for float values
     # expect:   false   false   false    false   false   false    true    false    true
-    x_corner = [1., 1., 1., np.nan, np.nan, np.nan, np.inf, np.inf, np.NINF]
-    y_corner = [np.nan, np.inf, np.NINF, np.nan, np.inf, np.NINF, np.inf, np.NINF, np.NINF]
+    x_corner = [1., 1., 1., np.nan, np.nan, np.nan, np.inf, np.inf, -np.inf]
+    y_corner = [np.nan, np.inf, -np.inf, np.nan, np.inf, -np.inf, np.inf, -np.inf, -np.inf]
 
     test_data_float16 = [
         pytest.param(

tests/requirements_tensorflow

@@ -1,15 +1,21 @@
-# tensorflow-intel inside tensorflow still requires numpy<2.0.0
-numpy==1.26.4
+# test ovc with NumPy 2.x on Ubuntu 24 with default Python 3.12
+# test against NumPy 1.x with older Python versions
+# tensorflow-intel 2.18.0 depends on numpy<2.1.0 and >=1.26.0
+numpy==1.26.4; python_version < "3.12"
+numpy==2.0.2; python_version >= "3.12"
 pytest==7.0.1
 pytest-xdist[psutil]==3.6.1
 pytest-html==4.1.1
 transformers==4.45.1
 # install exact keras version since tensorflow depends and has no upper bound for it
 keras==3.6.0
-tensorflow==2.17.0; platform_system != "Darwin" or platform_machine != "x86_64"
+tensorflow==2.18.0; python_version >= "3.12" and (platform_system != "Darwin" or platform_machine != "x86_64")
+tensorflow==2.17.0; python_version < "3.12" and (platform_system != "Darwin" or platform_machine != "x86_64")
 tensorflow==2.16.2; platform_system == "Darwin" and platform_machine == "x86_64"
+# install explicit version of wrapt to avoid "this __dict__ descriptor does not support '_DictWrapper' objects" error from TensorFlow 2.18
+wrapt==1.15.0; python_version >= "3.12"
 # tensorflow-text is not available for both Windows and ARM platforms
-tensorflow-text==2.17.0; platform_system == "Linux" and platform_machine == "x86_64"
+tensorflow-text==2.17.0; python_version < "3.12" and platform_system == "Linux" and platform_machine == "x86_64"
 tensorflow-hub==0.16.1
-jax==0.4.33
+jax==0.4.35
 defusedxml==0.7.1

tools/mo/requirements_tf.txt

@@ -1,6 +1,6 @@
 -c ../constraints.txt
 h5py
-tensorflow>=1.15.5,<2.18.0
+tensorflow>=1.15.5,<2.19.0
 numpy>=1.16.6,<1.27
 networkx
 defusedxml

tools/mo/requirements_tf2.txt

@@ -1,6 +1,6 @@
 -c ../constraints.txt
 h5py
-tensorflow>=2.5,<2.18.0
+tensorflow>=2.5,<2.19.0
 numpy>=1.16.6,<1.27
 networkx
 defusedxml