3 ms·
> model compatibility JAX will compile your model with XLA today, letting you run it on Nvidia, Apple Silicon (CPU for now), google's TPU, etc. Overall, XLA r
by nestorD 4y ago
> model compatibility
JAX will compile your model with XLA today, letting you run it on Nvidia, Apple Silicon (CPU for now), google's TPU, etc.
Overall, XLA really shines on TPUs, is doing very well on GPUs and sadly lags behind on CPUs (reaching about single core performance).
As a bonus, JAX is basically Numpy++: you can use it to build scientific/numerical software and run it easily on GPU and TPU!