>?JAX 是一個用于高性能數值計算的 Python 庫，特別為機器學習領域的高性能計算設計。它的 API 基于 Numpy 構建，包含豐富的數值計算與科學計算函數。JAX其實是 TensorFlow 的一個簡化庫，結合 Autograd 和 XLA，可以支持部分 TensorFlow 的功能，但是比 TensorFlow 更加簡潔易用。

>?Python 和 Numpy 的廣泛使用，使得?JAX 十分簡潔、靈活、易于上手，學習成本也比較低。除了 Numpy 的 API 外，JAX 還包含一系列可拓展、可組合的系統功能，有力地支持了機器學習研究。這些功能特性主要包括：

- 可差分：基于梯度的優化方法在機器學習領域具有十分重要的作用。JAX 可通過grad、hessian、jacfwd 和 jacrev 等函數轉換，原生支持任意數值函數的前向和反向模式的自動微分。

- 向量化：在機器學習中，通常需要在大規模的數據上運行相同的函數，例如計算整個批次的損失或每個樣本的損失等。JAX 通過 vmap 變換提供了自動矢量化算法，大大簡化了這種類型的計算，這使得研究人員在處理新算法時無需再去處理批量化的問題。JAX 同時還可以通過 pmap 轉換支持大規模的數據并行，從而優雅地將單個處理器無法處理的大數據進行處理。

- JIT編譯：XLA (Accelerated Linear Algebra, 加速線性代數) 被用于 JIT 即時編譯，在 GPU 和云 TPU 加速器上執行 JAX 程序。JIT 編譯與 JAX 的 API (與 Numpy 一致的數據函數) 為研發人員提供了便捷接入高性能計算的可能，無需特別的經驗就能將計算運行在多個加速器上。