本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。

收獲接近16.6萬個Star、見證深度學習崛起的TensorFlow，地位已岌岌可危。

并且這次，沖擊不是來自老對手PyTorch，而是自家新秀JAX。

最新一波AI圈熱議中，連fast.ai創始人Jeremy Howard都下場表示：

JAX正逐漸取代TensorFlow這件事，早已廣為人知了。現在它就在發生（至少在谷歌內部是這樣）。

LeCun更是認為，深度學習框架之間的激烈競爭，已經進入了一個新的階段。

LeCun表示，當初谷歌的TensorFlow確實比Torch更火。然而Meta的PyTorch出現之后，現在其受歡迎程度已經超過TensorFlow了。

現在，包括Google Brain、DeepMind以及不少外部項目，都已經開始用上JAX。

典型例子就是最近爆火的DALL·E Mini，為了充分利用TPU，作者采用了JAX進行編程。有人用過后感嘆：

這可比PyTorch快多了。

據《商業內幕》透露，預計在未來幾年內，JAX將覆蓋谷歌所有采用機器學習技術的產品。

這樣看來，如今大力在內部推廣JAX，更像是谷歌在框架上發起的一場“自救”。

JAX從何而來？

關于JAX，谷歌其實是有備而來。

早在2018年的時候，它就由谷歌大腦的一個三人小團隊給搭出來了。

研究成果發表在了題為Compiling machine learning programs via high-level tracing的論文中：

Jax是一個用于高性能數值計算的Python庫，而深度學習只是其中的功能之一。

自誕生以來，它受歡迎的程度就一直在上升。

最大的特點就是快。

一個例子感受一下。

比如求矩陣的前三次冪的和，用NumPy實現，計算需要約478毫秒。

用JAX就只需要5.54 毫秒，比NumPy快86倍。

為什么這么快？原因有很多，包括：

1、NumPy加速器。NumPy的重要性不用多說，用Python搞科學計算和機器學習，沒人離得開它，但它原生一直不支持GPU等硬件加速。

JAX的計算函數API則全部基于NumPy，可以讓模型很輕松在GPU和TPU上運行。這一點就拿捏住了很多人。

2、XLA。XLA（Accelerated Linear Algebra）就是加速線性代數，一個優化編譯器。JAX建立在XLA之上，大幅提高了JAX計算速度的上限。

3、JIT。研究人員可使用XLA將自己的函數轉換為實時編譯（JIT）版本，相當于通過向計算函數添加一個簡單的函數修飾符，就可以將計算速度提高幾個數量級。

除此之外，JAX與Autograd完全兼容，支持自動差分，通過grad、hessian、jacfwd和jacrev等函數轉換，支持反向模式和正向模式微分，并且兩者可以任意順序組成。

當然，JAX也是有一些缺點在身上的。

比如：

1、雖然JAX以加速器著稱，但它并沒有針對CPU計算中的每個操作進行充分優化。

2、JAX還太新，沒有形成像TensorFlow那樣完整的基礎生態。因此它還沒有被谷歌以成型產品的形式推出。

3、debug需要的時間和成本不確定，“副作用”也不完全明確。

4、不支持Windows系統，只能在上面的虛擬環境中運行。

5、沒有數據加載器，得借用TensorFlow或PyTorch的。

……

盡管如此，簡單、靈活又好用的JAX還是率先在DeepMind中流行起來。2020年誕生的一些深度學習庫Haiku和RLax等都是基于它開發。

這一年，PyTorch原作者之一Adam Paszke，也全職加入了JAX團隊。

目前，JAX的開源項目在GitHub上已有18.4k標星，比TensorFlow高了不少了。

值得注意的是，在此期間，有不少聲音都表示它很可能取代TensorFlow。

一方面是因為JAX的實力，另一方面主要還是跟TensorFlow自身的很多原因有關。

為什么谷歌要轉投JAX？

誕生于2015年的TensorFlow，曾經也風靡一時，推出后很快超過了Torch、Theano和Caffe等一眾“弄潮兒”，成為最受歡迎的機器學習框架。

然而在2017年，煥然一新的PyTorch“卷土重來”。

這是Meta基于Torch搭建的機器學習庫，由于上手簡單、通俗易懂，很快受到一眾研究者的青睞，甚至有超過TensorFlow的趨勢。

相比之下，TensorFlow卻在頻繁更新和界面迭代中變得越來越臃腫，逐漸失去了開發者的信任。

（從Stack Overflow上的提問占比來看，PyTorch逐年上升，TensorFlow卻一直停滯不前）

在競爭之中，TensorFlow的缺點逐漸暴露出來，API不穩定、實現復雜、學習成本高等問題并沒有隨著更新解決多少，反而結構變得更復雜了。

相比之下，TensorFlow卻沒有繼續發揮比較能打的“運行效率”等優勢。

在學術界，PyTorch的使用率正逐漸超過TensorFlow。

尤其是在各大頂會如ACL、ICLR中，使用PyTorch實現的算法框架近幾年已經占據了超過80%，相比之下TensorFlow的使用率還在不斷下降。

也正是因此，谷歌坐不住了，試圖用JAX奪回對機器學習框架的“主導權”。

雖然JAX名義上不是“專為深度學習構建的通用框架”，然而從發布之初起，谷歌的資源就一直在向JAX傾斜。

一方面，谷歌大腦和DeepMind逐漸將更多的庫構建在JAX上。

包括谷歌大腦的Trax、Flax、Jax-md，以及DeepMind的神經網絡庫Haiku和強化學習庫RLax等，都是基于JAX構建的。

據谷歌官方表示：

JAX生態系統開發中，也會考慮確保其與現有TensorFlow庫（如Sonnet和TRFL）的設計（盡可能）保持一致。

另一方面，更多的項目也開始基于JAX實現，最近爆火的DALL·E mini項目就是其中一個。

由于能更好地利用谷歌TPU的優勢，JAX在運行性能上比PyTorch要好得多，更多之前搭建在TensorFlow上的工業界項目也正在轉投JAX。

甚至有網友調侃JAX如今爆火的原因：可能是TensorFlow的使用者實在無法忍受這個框架了。

那么，JAX到底有沒有希望替代TensorFlow，成為與PyTorch抗衡的新力量呢？

更看好哪一個框架？

總體來看，很多人還是很堅定地站PyTorch。

他們似乎不喜歡谷歌每年都出一個新框架的速度。

“JAX雖然很吸引人，但還不夠具備“革命性”的能力促使大家拋棄PyTorch來使用它。”

但看好JAX的也并非少數。

就有人表示，PyTorch是很完美，但JAX也在縮小差距。

甚至還有人瘋狂給JAX打call，表示它比PyTorch要厲害10倍，并稱：如果Meta不繼續加把勁兒的話谷歌就會贏了。（手動狗頭）

不過，總有不怎么care誰輸誰贏的人，他們的目光放得很長遠：

沒有最好，只有更好。最重要的是更多玩家和好的idea統統都加入進來，讓開源和真正優秀的創新畫上等號。

項目地址：
https://github.com/google/jax