让 AI 替芯片优化程序，像请人改装赛车：如果每支队伍挑不同赛道计时，成绩就没法比较。JAXBench 补上了 TPU 领域的统一考场。TPU 是 Google 面向神经网络计算设计的专用芯片；内核则是芯片上执行矩阵乘法等基础操作的小程序，优化后可缩短模型训练或推理时间。
JAXBench 收录 50 个 JAX 工作负载，其中 17 个来自 Llama-3.1、DeepSeek-V3、Mamba-2 和 AlphaFold2 等架构，另有 33 个从 GPU 内核基准 KernelBench 转换而来。它们在 TPU v6e 上按统一规则检查结果是否正确并测量速度，8 项任务还配有人工优化版本作为参照。
最值得注意的是，作者发现：面对资料稀少的 Pallas——编写 TPU 底层程序的工具——提供整理过的 TPU 文档，比单纯扩大模型更重要。使用 Gemini 3 Flash 时，单次生成的正确率从 5.8% 升至 37.3%，并最终解出 50 项中的 48 项。换句话说，这个基准不只方便给 Agent 排名，也把瓶颈指向了更准确的硬件知识。