JAX

Compiling a Physics Engine into One Kernel: jit, grad and vmap in Practice

A case study of jax-spring-sim: forces derived by autodiff from a single energy function, a rollout compiled into one XLA kernel, and measured numbers for what compilation and reverse-mode differentiation actually cost.