JAX vs. numba

[rmojgani]

Does anyone know how the JAX-CPU version compares to basic numpy?
I assume the same?

If so, can we have

• the code structured to run without requiring JAX
• the code to go to numba JIT for parallel CPU

?

[jakharkaran]

Here's a discussion on JAX-CPU vs numpy: the conclusion is it depends.
https://jax.readthedocs.io/en/latest/faq.html#is-jax-faster-than-numpy

These would be good to have features, I have created a TODO file and added these.
https://github.com/envfluids/py2d/blob/main/TODO.md

[rmojgani]

I don't know if numba jit is going to be faster than JAX on cpu, so good to find some evidence before committing to it

[jakharkaran]

If we wish to use GPU it probably would be better to use Pytorch instead of Jax easing the coupling of NNs and considering the development efforts behind pytorch.

Someone, in the future, can do these comparisons if needed and take a look at specific libraries and if one suite better than the other for us.