2 papers
cs.DC2026
Orbax: Distributed Checkpointing with JAX
Colin Gaffney, Shutong Li, Daniel Ng +13
In a landscape of high-performance distributed ML systems, JAX has emerged as a framework of choice. However, JAX's modular design philosophy leaves it without a standardized check…
cs.SE2026
A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX
Stoyan Nikolov, Bernhard Konrad, Moritz Gronbach +5
The rapid development of AI-based products and their underlying models has led to constant innovation in deep learning frameworks. Google has been pioneering machine learning usage…