Speaker
Description
There is a growing interest in differentiable simulations that have fast execution time and yield additional gradient information. Most differentiable codes use a single autodiff or truncated power series algebra library and then implement standard optics and particle tracking logic on top. This approach can be limiting due to performance/compilation/ease of use tradeoffs, especially when wanting to use bespoke hardware accelerators. We present a new code for differentiable simulations, JACC (Jax for ACCelerators), that can create an interface for several libraries (Jax, PyTorch, NVIDIA Warp, finite differences) though dynamic beamline code synthesis. It can also stitch together pre-compiled or eager execution kernels. This helps with easy debugging while still enabling performance-critical kernel fusion, compilation, and distributed execution. An important application is end-to-end differentiable reinforcement learning (RL), where both Jax and PyTorch codebases are popular. Common Xsuite elements are implemented and a template for custom differentiable models provided. We present benchmarks including collective effects and apertures on CPUs, GPUs, and to our knowledge for the first time large clusters of Google TPUs. Suggested settings and practical recommendations are discussed for using JACC with digital twins, RL, and optimizers.
| In which format do you inted to submit your paper? | LaTeX |
|---|