5 results
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.
