?自 2018 年底推出以來(lái)，JAX 的受歡迎程度一直在穩(wěn)步提升。2020 年，DeepMind 宣布使用 JAX 來(lái)加速其研究。越來(lái)越多來(lái)自谷歌大腦（Google Brain）和其他機(jī)構(gòu)的項(xiàng)目也都在使用 JAX。 

目前，在 JAX 的 GitHub 項(xiàng)目主頁(yè)，Star 量已經(jīng)達(dá)到了 16.3k。?

項(xiàng)目地址：https://github.com/google/jaxJAX 是一個(gè)非常有前途的項(xiàng)目，并且用戶一直在穩(wěn)步增長(zhǎng)。JAX 已經(jīng)在深度學(xué)習(xí)、機(jī)器人 / 控制系統(tǒng)、貝葉斯方法和科學(xué)模擬等諸多領(lǐng)域得到了廣泛應(yīng)用。

如此，是否意味著 JAX 也將成為下一個(gè)大型深度學(xué)習(xí)框架？近日，發(fā)表在 AssemblyAI 博客上的文章《Why You Should (or Shouldn't) Be Using JAX in 2022》中，作者 Ryan O'Connor 為我們深入解讀了 JAX 的概念、使用 JAX 的理由以及是否應(yīng)該使用 JAX 等。

JAX 簡(jiǎn)介

JAX 不是一個(gè)深度學(xué)習(xí)框架或庫(kù)，其設(shè)計(jì)初衷也不是成為一個(gè)深度學(xué)習(xí)框架或庫(kù)。簡(jiǎn)而言之，JAX 是一個(gè)包含可組合函數(shù)轉(zhuǎn)換的數(shù)值計(jì)算庫(kù)。正如我們所看到的，深度學(xué)習(xí)只是 JAX 功能的一小部分：

?JAX 的定位科學(xué)計(jì)算（Scientific Computing）和函數(shù)轉(zhuǎn)換（Function Transformations）的交叉融合，具有除訓(xùn)練深度學(xué)習(xí)模型以外的一系列能力，包括如下：

• 即時(shí)編譯（Just-in-Time Compilation）
• 自動(dòng)并行化（Automatic Parallelization）
• 自動(dòng)向量化（Automatic Vectorization）
• 自動(dòng)微分（Automatic Differentiation）

使用 JAX 的原因有哪些？

簡(jiǎn)而言之，是速度。這是 JAX 與任何用例相關(guān)的一種通用能力。讓我們使用 NumPy 和 JAX 對(duì)矩陣的前三個(gè)冪求和（按元素）。 

首先是 NumPy 實(shí)現(xiàn)。我們發(fā)現(xiàn)，該計(jì)算大約需要 851 毫秒。?

然后使用 JAX 實(shí)現(xiàn)該計(jì)算：JAX 僅在 5.54 毫秒內(nèi)執(zhí)行完成該計(jì)算，速度是 NumPy 的 150 倍以上。

?JAX 的速度比 NumPy 快了 N 個(gè)數(shù)量級(jí)。需要注意，JAX 使用的是 TPU，NumPy 使用了 CPU，以此強(qiáng)調(diào) JAX 的速度上限遠(yuǎn)高于 NumPy。

作者列出了以下六條可能想要使用 JAX 的理由：

• NumPy 加速器。NumPy 是使用 Python 進(jìn)行科學(xué)計(jì)算的基礎(chǔ)包之一，但它僅與 CPU 兼容。JAX 提供了 NumPy 的實(shí)現(xiàn)（具有幾乎相同的 API），可以非常輕松地在 GPU 和 TPU 上運(yùn)行。對(duì)于許多用戶而言，僅此一項(xiàng)功能就足以證明使用 JAX 的合理性；
• XLA。XLA（Accelerated Linear Algebra）是專為線性代數(shù)設(shè)計(jì)的全程序優(yōu)化編譯器。JAX 建立在 XLA 之上，顯著提高了計(jì)算速度上限；
• JIT。JAX 允許用戶使用 XLA 將自己的函數(shù)轉(zhuǎn)換為即時(shí)編譯（JIT）版本。這意味著可以通過(guò)在計(jì)算函數(shù)中添加一個(gè)簡(jiǎn)單的函數(shù)裝飾器（decorator）來(lái)將計(jì)算速度提高幾個(gè)數(shù)量級(jí)；
• Auto-differentiation。JAX 將 Autograd（自動(dòng)區(qū)分原生 Python 代碼和 NumPy 代碼）和 XLA 結(jié)合在一起，它的自動(dòng)微分能力在科學(xué)計(jì)算的許多領(lǐng)域都至關(guān)重要。JAX 提供了幾個(gè)強(qiáng)大的自動(dòng)微分工具；
• 深度學(xué)習(xí)。雖然 JAX 本身不是深度學(xué)習(xí)框架，但它的確為深度學(xué)習(xí)提供了一個(gè)很好的基礎(chǔ)。很多構(gòu)建在 JAX 之上的庫(kù)旨在提供深度學(xué)習(xí)功能，包括 Flax、Haiku 和 Elegy。甚至在最近的一些 PyTorch 與 TensorFlow 文章中強(qiáng)調(diào)了 JAX 作為一個(gè)值得關(guān)注的「框架」，并推薦其用于基于 TPU 的深度學(xué)習(xí)研究。JAX 對(duì) Hessians 的高效計(jì)算也與深度學(xué)習(xí)相關(guān)，因?yàn)樗鼈兪垢唠A優(yōu)化技術(shù)更加可行；
• 通用可微分編程范式（General Differentiable Programming Paradigm ）。雖然我們可以使用 JAX 來(lái)構(gòu)建和訓(xùn)練深度學(xué)習(xí)模型，但它也為通用可微編程提供了一個(gè)框架。這意味著 JAX 可以通過(guò)使用基于模型的機(jī)器學(xué)習(xí)方法來(lái)解決問(wèn)題，從而可以利用數(shù)十年研究建立起的給定領(lǐng)域的先驗(yàn)知識(shí)。?

JAX 轉(zhuǎn)換

?到目前為止，我們已經(jīng)討論了 XLA 以及它如何允許 JAX 在加速器上實(shí)現(xiàn) NumPy；但請(qǐng)記住，這只是 JAX 定義的一半。JAX 不僅為強(qiáng)大的科學(xué)計(jì)算提供了工具，而且還為可組合的函數(shù)轉(zhuǎn)換提供了工具。

舉例來(lái)說(shuō)如果我們對(duì)標(biāo)量值函數(shù) f(x) 使用梯度函數(shù)轉(zhuǎn)換，那么我們將得到一個(gè)向量值函數(shù) f'(x)，它給出了函數(shù)在 f(x) 域中任意點(diǎn)的梯度。?

在函數(shù)上使用 grad() 可以讓我們得到域中任意點(diǎn)的梯度

?JAX 包含了一個(gè)可擴(kuò)展系統(tǒng)來(lái)實(shí)現(xiàn)這樣的函數(shù)轉(zhuǎn)換，有四種典型方式：

• Grad() 進(jìn)行自動(dòng)微分；
• Vmap() 自動(dòng)向量化；
• Pmap() 并行化計(jì)算；
• Jit() 將函數(shù)轉(zhuǎn)換為即時(shí)編譯版本。

使用 grad() 進(jìn)行自動(dòng)微分

訓(xùn)練機(jī)器學(xué)習(xí)模型需要反向傳播。在 JAX 中，就像在 Autograd 中一樣，用戶可以使用 grad() 函數(shù)來(lái)計(jì)算梯度。

舉例來(lái)說(shuō)，如下是對(duì)函數(shù) f(x) = abs(x^3) 求導(dǎo)。我們可以看到，當(dāng)求 x=2 和 x=-3 處的函數(shù)及其導(dǎo)數(shù)時(shí)，我們得到了預(yù)期的結(jié)果。?

那么 grad() 能微分到什么程度？JAX 通過(guò)重復(fù)應(yīng)用 grad() 使得微分變得很容易，如下程序我們可以看到，輸出函數(shù)的三階導(dǎo)數(shù)給出了 f'''(x)=6 的恒定預(yù)期輸出。

可能有人會(huì)問(wèn)，grad() 可以用在哪些方面？標(biāo)量值函數(shù)：grad() 采用標(biāo)量值函數(shù)的梯度，將標(biāo)量 / 向量映射到標(biāo)量函數(shù)。此外還有向量值函數(shù)：對(duì)于將向量映射到向量的向量值函數(shù)，梯度的類似物是雅可比矩陣。使用 jacfwd() 和 jacrev()，JAX 返回一個(gè)函數(shù)，該函數(shù)在域中的某個(gè)點(diǎn)求值時(shí)產(chǎn)生雅可比矩陣。

?從深度學(xué)習(xí)角度來(lái)看，JAX 使得計(jì)算 Hessians 變得非常簡(jiǎn)單和高效。由于 XLA，JAX 可以比 PyTorch 更快地計(jì)算 Hessians，這使得實(shí)現(xiàn)諸如 AdaHessian 這樣的高階優(yōu)化更加快速。

下面代碼是在 PyTorch 中對(duì)一個(gè)簡(jiǎn)單的輸入總和進(jìn)行 Hessian：?

正如我們所看到的，上述計(jì)算大約需要 16.3 ms，在 JAX 中嘗試相同的計(jì)算：

使用 JAX，計(jì)算僅需 1.55 毫秒，比 PyTorch 快 10 倍以上：JAX 可以非常快速?地計(jì)算 Hessians，使得高階優(yōu)化更加可行。

使用 vmap() 自動(dòng)向量化

JAX 在其 API 中還有另一種變換：vmap() 自動(dòng)向量化。以下是矢量化向量加法展示：?

使用 pmap() 實(shí)現(xiàn)自動(dòng)并行化

分布式計(jì)算變得越來(lái)越重要，在深度學(xué)習(xí)中尤其如此，如下圖所示，SOTA 模型已經(jīng)發(fā)展到超大規(guī)模。

?得益于 XLA，JAX 可以輕松地在加速器上進(jìn)行計(jì)算，但 JAX 也可以輕松地使用多個(gè)加速器進(jìn)行計(jì)算，即使用單個(gè)命令 - pmap() 執(zhí)行 SPMD 程序的分布式訓(xùn)練。

我們以向量矩陣乘法為例，如下為非并行向量矩陣乘法：?

使用 JAX，我們可以輕松地將這些計(jì)算分布在 4 個(gè) TPU 上，只需將操作包裝在 pmap() 中即可。這允許用戶在每個(gè) TPU 上同時(shí)執(zhí)行一個(gè)點(diǎn)積，顯著提高了計(jì)算速度（對(duì)于大型計(jì)算而言）。

使用 jit() 加快功能

JIT 編譯是一種執(zhí)行代碼的方法，介于解釋（interpretation）和 AoT（ahead-of-time）編譯之間。重要的是，JIT 編譯器在運(yùn)行時(shí)將代碼編譯成快速的可執(zhí)行文件，但代價(jià)是首次運(yùn)行速度較慢。

JIT 不是一次將一個(gè)操作分配給 GPU 內(nèi)核，而是使用 XLA 將一系列操作編譯成一個(gè)內(nèi)核，從而為函數(shù)提供端到端編譯的高效 XLA 實(shí)現(xiàn)。

以下圖為例，代碼定義了一個(gè)函數(shù)：用三種方式計(jì)算 5000 x 5000 矩陣——一次使用 NumPy，一次使用 JAX，還有一次在 JIT 編譯的函數(shù)版本上使用 JAX。我們首先在 CPU 上進(jìn)行實(shí)驗(yàn)：?

JAX 對(duì)于逐元素計(jì)算明顯更快，尤其是在使用 jit 時(shí)。

我們看到 JAX 比 NumPy 快 2.3 倍以上，當(dāng)我們 JIT 函數(shù)時(shí)，JAX 比 NumPy 快 30 倍。這些結(jié)果已經(jīng)令人印象深刻，但讓我們繼續(xù)看，讓 JAX 在 TPU 上進(jìn)行計(jì)算：

?當(dāng) JAX 在 TPU 上執(zhí)行相同的計(jì)算時(shí)，它的相對(duì)性能會(huì)進(jìn)一步提升（NumPy 計(jì)算仍在 CPU 上執(zhí)行，因?yàn)樗恢С?TPU 計(jì)算）在這種情況下，我們可以看到 JAX 比 NumPy 快了驚人的 13 倍，如果我們同時(shí)在 TPU 上 JIT 函數(shù)和計(jì)算，我們會(huì)發(fā)現(xiàn) JAX 比 NumPy 快 80 倍。

當(dāng)然，這種速度的大幅提升是有代價(jià)的。JAX 對(duì) JIT 允許的函數(shù)進(jìn)行了限制，盡管通常允許僅涉及上述 NumPy 操作的函數(shù)。此外，通過(guò) Python 控制流進(jìn)行 JIT 處理存在一些限制，因此在編寫函數(shù)時(shí)須牢記這一點(diǎn)。

2022 年了，我該用 JAX 嗎？

很遺憾，這個(gè)問(wèn)題的答案還是「視情況而定」。是否遷移到 JAX 取決于你的情況和目標(biāo)。為具體分析是否應(yīng)該（或不應(yīng)該）在 2022 年使用 JAX，這里將建議匯總到下面的流程圖中，并針對(duì)不同的興趣領(lǐng)域提供不同的圖表。?

科學(xué)計(jì)算

?如果你對(duì) JAX 在通用計(jì)算感興趣，首先要問(wèn)的問(wèn)題就是——是否只嘗試在加速器上運(yùn)行 NumPy？如果答案是肯定的，那么你顯然應(yīng)該開(kāi)始遷移到 JAX。

如果你不只處理數(shù)字而是參與動(dòng)態(tài)計(jì)算建模，那么是否應(yīng)該使用 JAX 將取決于具體用例。如果大部分工作是在 Python 中使用大量自定義代碼完成的，那么開(kāi)始學(xué)習(xí) JAX 以增強(qiáng)工作流程是值得的。

如果大部分工作不在 Python 中，但你想構(gòu)建的是某種基于模型 / 神經(jīng)網(wǎng)絡(luò)的混合系統(tǒng)，那么使用 JAX 可能是值得的。

如果大部分工作不使用 Python，或者你正在使用一些專門的軟件進(jìn)行研究（熱力學(xué)、半導(dǎo)體等），那么 JAX 可能是不合適的工具，除非你想從這些程序中導(dǎo)出數(shù)據(jù)，用來(lái)做自定義計(jì)算。如果你感興趣的領(lǐng)域更接近物理 / 數(shù)學(xué)并包含計(jì)算方法（動(dòng)力系統(tǒng)、微分幾何、統(tǒng)計(jì)物理）并且大部分工作都在例如 Mathematica 上，那么堅(jiān)持使用目前的工具才是值得的，特別是在已有大型自定義代碼庫(kù)的情形下。?

深度學(xué)習(xí)

?雖然我們已經(jīng)強(qiáng)調(diào)過(guò)，JAX 不是專為深度學(xué)習(xí)構(gòu)建的通用框架，但 JAX 速度很快且具有自動(dòng)微分功能，你肯定想知道使用 JAX 進(jìn)行深度學(xué)習(xí)是什么樣的。

若想在 TPU 上進(jìn)行訓(xùn)練，那么你應(yīng)該開(kāi)始使用 JAX，尤其是如果當(dāng)前正在使用的是 PyTorch。雖然有 PyTorch-XLA 存在，但使用 JAX 進(jìn)行 TPU 訓(xùn)練絕對(duì)是更好的體驗(yàn)。如果你正在研究的是「非標(biāo)準(zhǔn)」架構(gòu) / 建模，例如 SDE-Nets，那么也絕對(duì)應(yīng)該嘗試一下 JAX。此外，如果你想利用高階優(yōu)化技術(shù)，JAX 也是要嘗試的東西。

如果你不是在構(gòu)建特殊的架構(gòu)，只是在 GPU 上訓(xùn)練常見(jiàn)的架構(gòu)，那么你現(xiàn)在可能應(yīng)該堅(jiān)持使用 PyTorch 或 TensorFlow。然而，這個(gè)建議可能會(huì)在未來(lái)一兩年內(nèi)快速發(fā)生變化。雖然 PyTorch 仍然在研究領(lǐng)域占據(jù)主導(dǎo)地位，但使用 JAX 的論文數(shù)量一直在穩(wěn)步增長(zhǎng)。隨著 DeepMind 和谷歌重量級(jí)玩家不斷開(kāi)發(fā)用于 JAX 的高級(jí)深度學(xué)習(xí) API，在幾年內(nèi) JAX 可能會(huì)出現(xiàn)爆炸性的增長(zhǎng)率。

這意味著你至少應(yīng)該稍微熟悉一下 JAX，如果你是研究人員的話更應(yīng)如此。?

深度學(xué)習(xí)初學(xué)者

但如果我只是個(gè)初學(xué)者呢？情況會(huì)有些不一樣。

如果你有興趣了解深度學(xué)習(xí)并實(shí)現(xiàn)一些想法，你應(yīng)該使用 JAX 或 PyTorch。如果你想自上而下學(xué)習(xí)深度學(xué)習(xí)，或有一些 Python 軟件的經(jīng)驗(yàn)，則應(yīng)該從 PyTorch 入手。如果你想自下而上地學(xué)習(xí)深度學(xué)習(xí)，或具有數(shù)學(xué)背景，你可能會(huì)發(fā)現(xiàn) JAX 很直觀。在這種情況下，在進(jìn)行任何大型項(xiàng)目之前，請(qǐng)確保了解如何使用 JAX。

如果你對(duì)深度學(xué)習(xí)感興趣，又想轉(zhuǎn)行相關(guān)的職位，那么你需要使用 PyTorch 或 TensorFlow。盡管最好是同時(shí)熟悉兩個(gè)框架，但你必須知道 TensorFlow 被普遍認(rèn)為是「行業(yè)」框架，不同框架的職位發(fā)布數(shù)量證明了這一點(diǎn)：

如果你是一個(gè)沒(méi)有數(shù)學(xué)或軟件背景但想學(xué)習(xí)深度學(xué)習(xí)的初學(xué)者，那么你不會(huì)想使用 JAX。相反，Keras 是更好的選擇。

不該使用 JAX 的四條理由

?雖然上文已經(jīng)討論了很多 JAX 的正面反饋，它有潛力極大地提升用戶程序的性能。但作者同時(shí)列舉了以下四條不該使用 JAX 的理由：

• ?JAX 仍然被官方認(rèn)為是一個(gè)實(shí)驗(yàn)性框架。JAX 是一個(gè)相對(duì)「年輕」的項(xiàng)目。目前，JAX 仍被視為一個(gè)研究項(xiàng)目，而不是成熟的谷歌產(chǎn)品，因此如果用戶正在考慮遷移到 JAX，請(qǐng)記住這一點(diǎn)；
• 使用 JAX 一定要勤勉。調(diào)試的時(shí)間成本，或者更嚴(yán)重的是，未跟蹤副作用（untracked side effects）的風(fēng)險(xiǎn)可能導(dǎo)致那些沒(méi)有扎實(shí)掌握函數(shù)式編程的用戶不適用 JAX。在開(kāi)始將它用于正式項(xiàng)目之前，請(qǐng)確保自己了解使用 JAX 的常見(jiàn)缺陷；
• JAX 沒(méi)有針對(duì) CPU 計(jì)算進(jìn)行優(yōu)化。鑒于 JAX 是以「加速器優(yōu)先」的方式開(kāi)發(fā)的，因此每個(gè)操作的分派并未針對(duì) JAX 進(jìn)行完全優(yōu)化。在某些情況下，NumPy 實(shí)際上可能比 JAX 更快，尤其是對(duì)于小型程序而言，這是因?yàn)?JAX 引入了開(kāi)銷；
• JAX 與 Windows 不兼容。目前在 Windows 上不支持 JAX。如果用戶使用 Windows 系統(tǒng)但仍想嘗試 JAX，可以使用 Colab 或?qū)⑵浒惭b在虛擬機(jī)（VM）上。?