JAX 是 Google 和 NVIDIA 聯合開發的高性能數值計算庫，這兩年 JAX 生態快速發展，周邊工具鏈也日益完善了。如果你用過 NumPy 或 PyTorch，但還沒接觸過 JAX，這篇文章能幫助你快速上手。 圍繞 JAX 已經涌現出一批好用的庫：Flax 用來搭神經網絡，Optax 處理梯度和優化，Equinox 提供類似 PyTorch 的接口，Haiku 則是簡潔的函數式 API，
目前來看Google 是唯一一家在 AI 價值鏈上實現端到端垂直整合的公司。從基礎模型 (Gemini)、應用層 (ImageFX, Search with Gemini, NotebookLM)，到雲架構 (Google Cloud, Vertex AI) 以及硬件 (TPUs)，幾乎全都有所佈局。 長期以來Google 一直在通過提升自身能力來減少對 NVIDIA GPU 的依賴。這種
很多人剛接觸JAX都會有點懵——參數為啥要單獨傳？隨機數還要自己管key？這跟PyTorch的畫風完全不一樣啊。 其實根本原因就一個：JAX是函數式編程而不是面向對象那套，想明白這點很多設計就都説得通了。 先説個核心區別 PyTorch裏，模型是個對象，權重藏在裏面，訓練的時候自己更新自己。這是典型的面向對象思路，狀態封裝在對象內部。 JAX的思路完全反過來。模型定義是模型定義，參數是參數，兩邊分
TPU 訓練的真實效率往往取決於兩個核心要素：Shape 的穩定性與算子的融合度。 很多時候，JAX 任務之所以出現嚴重的性能瓶頸，並非算法本身設計有問題，而是忽視了 XLA 編譯器與底層硬件對“確定性”的極度偏好。基於大量實戰調優經驗，本文總結了八條能讓 JAX 訓練任務從“甚至跑不通”蜕變為“跑滿 TPU 算力”的工程經驗。 1、儘早鎖定 Shape TPU 喜歡靜態 Shape，JA
JAX生態這兩年在LLM訓練這塊追趕得挺快。PyTorch雖然還是主流但JAX在並行計算、TPU加速和API組合性上確實有些獨特的優勢。Google今天放出了Tunix這個庫，專門做LLM的後訓練——微調、強化學習、知識蒸餾這些都能搞。 Tunix是什麼 這是個構建在JAX之上的後訓練庫，和Flax NNX集成得比較緊密。主要解決三類問題： 監督微調（Supervised Fine-T
