6 results
JULY 14, 2026
To serve the 397B-parameter Qwen 3.5 Mixture-of-Experts (MoE) model on Ironwood TPUs, engineers developed a modular JAX/Pallas optimization stack that achieved up to a 4.7x inference speedup for prefill-heavy workloads. The team bypassed severe hardware sharding constraints by deploying a hybrid Data Parallelism and Expert Parallelism (DP+EP) topology, paired with custom low-level communication fusions like a hierarchical reduce-scatter to optimize cross-device token routing. Finally, by executing hardware-aware custom kernels—such as Batched Ragged Page Attention and a fully-fused Gated DeltaNet (GDN) block—they successfully saturated HBM bandwidth and TensorCore MXUs to push system throughput near its theoretical roofline limits.
AUG. 12, 2025 / Kaggle
Build and train a GPT2 model from scratch using JAX on Google TPUs, with a complete Python notebook for free-tier Colab or Kaggle. Learn how to define a hardware mesh, partition model parameters and input data for data parallelism, and optimize the model training process.
JULY 29, 2025 / AI
Max's journey introduces LQRax, a JAX-native LQR solver, which exemplifies the growing JAX robotics ecosystem that includes tools like Brax, MJX, and JaxSim, highlighting the benefits of JAX for computational efficiency in optimal control and simulation, and for seamlessly integrating model-based and learning-based approaches.
JULY 16, 2025 / Cloud
The Marin project aims to expand the definition of 'open' in AI to include the entire scientific process, not just the model itself, by making the complete development journey accessible and reproducible. This effort, powered by the JAX framework and its Levanter tool, allows for deep scrutiny, trust in, and building upon foundation models, fostering a more transparent future for AI research.
JUNE 24, 2025 / Kaggle
KerasHub enables users to mix and match model architectures and weights across different machine learning frameworks, allowing checkpoints from sources like Hugging Face Hub (including those created with PyTorch) to be loaded into Keras models for use with JAX, PyTorch, or TensorFlow. This flexibility means you can leverage a vast array of community fine-tuned models while maintaining full control over your chosen backend framework.
MAY 13, 2025 / TensorFlow
Keras Recommenders (KerasRS) is a new library announced to help developers build recommendation systems using APIs with building blocks for ranking and retrieval, and it can be installed via pip with support for JAX, TensorFlow, or PyTorch backends.
