add jax support

[Oisin-M]

Description

Contributor Declaration

By opening this pull request, I affirm the following:

• All authors agree to the Contributor License Agreement.
• The code follows the project's coding standards.
• I have performed self-review and added comments where needed.
• I have added or updated tests to verify that my changes are effective and functional.
• I have run all existing tests and confirmed they pass.