Bershatsky, D. (2024). YAX: JAX/FLAX Module Tracing, Evaluation, and Mutation (Version 0.0.0) [Computer software]. https://github.com/daskol/yax