[Features][Machine Learning] Jax support

[okbalefthanded]

• JAX support is currently unstable on windows, this feature is postponed until a stable release will be available.