[JAX]: Support jax.lax.select_n operation for JAX

[11happy]

Overview:
This pull request fixes #26570.

Testing:

• Tested the updated code.
• Verified that other functionalities remain unaffected.
  

Dependencies:

• No dependencies on other pull requests.

CC:

• @rkazants

[rkazants]

build_jenkins

[11happy]

humble ping! @rkazants

[11happy]

Hello @rkazants, please let me know if there's anything needed from my end.

[rkazants]

build_jenkins

[rkazants]

build_jenkins

[rkazants]

@11happy, please pay attention to failing CI jobs (code-style, build, etc.), it needs to be fixed.

[rkazants]

let use case_num starting from 2. Value 1 is impractical.

[11happy]

done

[rkazants]

build_jenkins

[rkazants]

This time I ran clang-format on my code what else I can do to ensure it passes CI jobs?. Thank you

Let CI complete and check failing jobs in GHA. You may see build issues that need to be fixed.

[rkazants]

build_jenkins

[rkazants]

@11happy, looks good. Let us see CI results. If it is fine, I will merge it.

Best regards,
Roman

[rkazants]

@11happy, please fix CI errors. Check GHA jobs. Example:

src/frontends/jax/src/CMakeFiles/openvino_jax_frontend.dir/op/select_n.cpp.o -c /__w/openvino/openvino/openvino/src/frontends/jax/src/op/select_n.cpp
/__w/openvino/openvino/openvino/src/frontends/jax/src/op/select_n.cpp:5:50: error: extra tokens at end of #include directive [-Werror,-Wextra-tokens]
#include "openvino/frontend/jax/node_context.hpp";
                                                 ^
                                                 //
/__w/openvino/openvino/openvino/src/frontends/jax/src/op/select_n.cpp:11:21: error: extra tokens at end of #include directive [-Werror,-Wextra-tokens]
#include "utils.hpp";
                    ^
                    //
2 errors generated.

[11happy]

@11happy, please fix CI errors. Check GHA jobs. Example:

src/frontends/jax/src/CMakeFiles/openvino_jax_frontend.dir/op/select_n.cpp.o -c /__w/openvino/openvino/openvino/src/frontends/jax/src/op/select_n.cpp
/__w/openvino/openvino/openvino/src/frontends/jax/src/op/select_n.cpp:5:50: error: extra tokens at end of #include directive [-Werror,-Wextra-tokens]
#include "openvino/frontend/jax/node_context.hpp";
                                                 ^
                                                 //
/__w/openvino/openvino/openvino/src/frontends/jax/src/op/select_n.cpp:11:21: error: extra tokens at end of #include directive [-Werror,-Wextra-tokens]
#include "utils.hpp";
                    ^
                    //
2 errors generated.

I have removed these extra semi colons.

[rkazants]

build_jenkins

[11happy]

I saw this gha fail:

numpy.int64'> - input_shape:[2, 3] ] - AssertionError: number of outputs not equal, 1 != 2
FAILED install/tests/layer_tests/jax_tests/test_select_n.py::TestSelectN::test_select_n[ ie_device:CPU - precision:FP16 - case_num:3 - input_type:<class 'numpy.int64'> - input_shape:[4, 5, 6] ] - AssertionError: number of outputs not equal, 1 != 2
FAILED install/tests/layer_tests/jax_tests/test_select_n.py::TestSelectN::test_select_n[ ie_device:CPU - precision:FP16 - case_num:3 - input_type:<class 'numpy.int64'> - input_shape:[7, 8, 9, 10] ] - AssertionError: number of outputs not equal, 1 != 2
=============== 10 failed, 1500 passed, 1510 warnings in 37.14s ================

How can I resolve this?

[rkazants]

build_jenkins

[rkazants]

How can I resolve this?

It looks a problem of JAX utils itself that helps to convert JAX model to TF model. I switched off a path via TF FE that is deprecated anyway.

Best regards,
Roman

[rkazants]

build_jenkins

[rkazants]

build_jenkins