[JAX FE] : Implement jax.lax.iota operation

[11happy]

Overview:

• This pull request fixes [Good First Issue][JAX FE]: Support jax.lax.iota operation for JAX #26569

Testing:

• TODO: FIX, Currently I am facing some assertion issues while implementing the tests.

CC:

• @rkazants

[11happy]

Hello @rkazants could you please help me with the tests.
Thank you

[rkazants]

build_jenkins

[rkazants]

@11happy, please fix failing CI jobs

[rkazants]

build_jenkins

[rkazants]

Hi @11happy, please fix the build:

/__w/openvino/openvino/openvino/bin/aarch64/Release/libopenvino_util.a  -ldl  vcpkg_installed/arm64-android/lib/libpugixml.a  -latomic -lm && :
ld: error: undefined symbol: ov::element::Type ov::frontend::jax::NodeContext::const_named_param<ov::element::Type>(std::__ndk1::basic_string<char, std::__ndk1::char_traits<char>, std::__ndk1::allocator<char> > const&) const
>>> referenced by iota.cpp:19 (src/frontends/jax/src/op/iota.cpp:19)
>>>               src/frontends/jax/src/CMakeFiles/openvino_jax_frontend.dir/op/iota.cpp.o:(ov::frontend::jax::op::translate_iota(ov::frontend::jax::NodeContext const&))
clang++: error: linker command failed with exit code 1 (use -v to see invocation)

There is no instantiation for const_named_param with ov::element::Type type and you need to define it.

Best regards,
Roman

[11happy]

I have defined it, however I am sure if this is correct way.

I am assuming this creates an uninitialized constant & we can get type as get_element_type

[11happy]

also I wanted to ask do we have any pre-commit hooks or something like that , as most of the time I push some changes build fails due to some formatting error.

[11happy]

@rkazants humble ping!
Thank you

[11happy]

@rkazants I think we don't need to define const_named_param as there is this convert_dtype I referenced from argmax implementation & its working . All tests are passing on my end.

[rkazants]

build_jenkins

[11happy]

I think here as both are used
via const_named_param we won't be needing this inputs size check.
there is this failing CI, currently removed this inputs size check.

thank you

[11happy]

Hi @rkazants,
I have updated the implementation, and it’s working. However, I’m not entirely sure which approach is correct to use both inputs via parameters or this way. Could you please clarify the preferred way to proceed?
Thank you!

[11happy]

humble ping!

[github-actions]

This PR will be closed in a week because of 2 weeks of no activity.

[11happy]

Humble Ping!
Hii @rkazants could please take a look as all tests are passing as i checked locally, is there any more changes/inputs required from my side.
Thank you

[p-wysocki]

@rkazants could you please take a look?

[11happy]

Hello :) is there any further input required from my side, I would like this PR to be completed.
Thank you

[github-actions]

This PR will be closed in a week because of 2 weeks of no activity.

[11happy]

gentle ping! is there any more input required from my side @rkazants to complete this PR.
Thank you

[mlukasze]

build_jenkins