Google Releases Gemini-Based AI Agent Tools to Accelerate TPU Development

Google has released Accelerator Agents, a suite of AI tools designed to accelerate machine learning development on Google Cloud TPU. This project leverages specialized Gemini-based AI agents to automate code migration, kernel optimization, and performance tuning.

Key Components

The project consists of two core agents.

• MaxCode: Converts existing PyTorch models and codebases to JAX. It generates code compatible with the MaxText framework and adopts a Human-in-the-Loop approach where developers review drafts.
• MaxKernel: Supports high-performance kernel development for TPUs. It assists in writing and debugging Pallas, JAX's kernel language, and provides optimization suggestions for porting CUDA/GPU kernels and improving MFU (Model FLOPs Utilization).

Usage Environment and Requirements

Execution requires a Google Cloud VM and a Python 3.11+ environment. MaxCode can run on CPU-only VMs, but MaxKernel recommends a TPU VM. Access to the Gemini API is essential for agent inference.

License and Status

This project follows the Apache License 2.0 and is currently in active development. However, it is not an officially supported Google product and is not included in the Google Open Source Software Vulnerability Rewards Program.