6 ms·
Jax uses the XLA compiler which is compatible with GPU and CPU.
by drdirk 2y ago
Jax uses the XLA compiler which is compatible with GPU and CPU.
- ein0p 2y agoYes, but its performance on GPU leaves much to be desired, and 20 times as much research comes out on PyTorch. Would you rather just build on that or laboriously port and debug the models and their weights, losses, dataset readers, training regimes etc etc?