Parallelizing neural networks on one GPU with JAX
willwhitney.com