JAX does JIT compilation which is not done by Cupy so you get better performance by fusing operations and compiling them cleverly ([0] is a good benchmark to ge

JAX does JIT compilation which is not done by Cupy so you get better performance by fusing operations and compiling them cleverly ([0] is a good benchmark to get a feeling for the performance of the various options). However, it is a trade-off, the jitting time will increase your run time (you can expect at least 10 seconds of compiling time) and JAX does not let you do in-place operations (however, it has workarounds). If you can afford the time needed to jit your code then I highly recommend JAX. [0]: https://github.com/dionhaefner/pyhpc-benchmarks https://github.com/dionhaefner/pyhpc-benchmarks