The appropriate audience for this content is Machine Learning Engineers (MLEs) and Systems Architects tasked with scaling deep learning workflows, particularly those evaluating whether to migrate from PyTorch or legacy TensorFlow to JAX/OpenXLA.
Abstract:
This presentation from the Fall 2025 JAX/OpenXLA DevLab details the evolution of Keras 3 as a high-level, multi-backend interface for JAX. The Keras team demonstrates how the framework handles the inherent statefulness of Keras layers within JAX's stateless functional paradigm, allowing for standard transformations like jax.jit, grad, and vmap.
The talk highlights three major pillars: high-level API simplicity for rapid prototyping, seamless interoperability with the new Flax NNX ecosystem, and specialized libraries for large-scale production. Specifically, Keras Hub is presented as a standardized gateway for SOTA models like Gemma and Stable Diffusion 3, while Keras RS introduces significant performance optimizations for recommender systems via "SparseCore" sharding on V5/V6 TPUs. As far as I can tell, the primary value proposition is "progressive disclosure of complexity"—enabling users to start with simple fit() calls and descend into custom JAX training loops or distributed hardware-specific optimizations as needed.
Keras: Deep Learning Framework for JAX | DevLab Fall 2025 Summary
0:04 Keras 3 Core Philosophy: The team emphasizes "progressive disclosure of complexity." The goal is to keep code elegant for 90% of use cases while allowing low-level JAX access for the remaining 10%.
5:32 JAX Backend Integration: Keras 3 converts stateful layers into stateless functions. This allows Keras models to be treated as standard JAX functions, compatible with jax.jit and jax.pmap.
11:31 Flax NNX Interoperability: A major update includes a dedicated NNX backend. Keras layers can now be used directly as Flax NNX modules, which is a necessary pivot given JAX's move toward the NNX state management model.
13:25 Customizing Training Loops: While model.fit() is the default, the team demonstrates overriding train_step and compute_loss. I don’t care for the boilerplate, but it's essential for those implementing non-standard gradients or GAN-style training.
19:33 Keras Hub & Pre-trained Models: Keras Hub acts as an "editorialized" version of Hugging Face. It focuses on SOTA models (Gemma, Llama, Stable Diffusion) with built-in tokenization and LoRA support.
27:34 Stable Diffusion 3 Case Study: The presenters demonstrate walking the text embedding space by interpolating between prompts (e.g., "dog" to "tentacled creature"). It’s a visual gimmick, but it proves the smooth manifold of the embedding space and Keras’s ability to handle complex diffusion backbones.
34:02 Keras RS (Recommender Systems): This is the most substantial part for enterprise users. Keras RS fills a gap in the JAX ecosystem for ranking and retrieval models, traditionally dominated by NVIDIA Merlin or TensorFlow RS.
42:47 Distributed Embeddings on TPU: The framework supports sharding massive embedding tables across "SparseCore" processors on V5/V6 TPUs. This is likely the only way to handle trillion-parameter recommender models efficiently.
45:24 Takeaways: The ecosystem is now modular. You can plug a Keras Hub transformer into a Keras RS ranking model and train it with a custom JAX loop. Reliability seems high, though I'm skeptical of how "simple" the TPU sharding remains once you leave the provided examples.