本文經(jīng)AI新媒體量子位（公眾號(hào)ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

收獲接近16.6萬(wàn)個(gè)Star、見(jiàn)證深度學(xué)習(xí)崛起的TensorFlow，地位已岌岌可危。

并且這次，沖擊不是來(lái)自老對(duì)手PyTorch，而是自家新秀JAX。

最新一波AI圈熱議中，連fast.ai創(chuàng)始人Jeremy Howard都下場(chǎng)表示：

JAX正逐漸取代TensorFlow這件事，早已廣為人知了。現(xiàn)在它就在發(fā)生（至少在谷歌內(nèi)部是這樣）。

LeCun更是認(rèn)為，深度學(xué)習(xí)框架之間的激烈競(jìng)爭(zhēng)，已經(jīng)進(jìn)入了一個(gè)新的階段。

LeCun表示，當(dāng)初谷歌的TensorFlow確實(shí)比Torch更火。然而Meta的PyTorch出現(xiàn)之后，現(xiàn)在其受歡迎程度已經(jīng)超過(guò)TensorFlow了。

現(xiàn)在，包括Google Brain、DeepMind以及不少外部項(xiàng)目，都已經(jīng)開(kāi)始用上JAX。

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

這可比PyTorch快多了。

據(jù)《商業(yè)內(nèi)幕》透露，預(yù)計(jì)在未來(lái)幾年內(nèi)，JAX將覆蓋谷歌所有采用機(jī)器學(xué)習(xí)技術(shù)的產(chǎn)品。

這樣看來(lái)，如今大力在內(nèi)部推廣JAX，更像是谷歌在框架上發(fā)起的一場(chǎng)“自救”。

JAX從何而來(lái)？

關(guān)于JAX，谷歌其實(shí)是有備而來(lái)。

早在2018年的時(shí)候，它就由谷歌大腦的一個(gè)三人小團(tuán)隊(duì)給搭出來(lái)了。

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

Jax是一個(gè)用于高性能數(shù)值計(jì)算的Python庫(kù)，而深度學(xué)習(xí)只是其中的功能之一。

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

最大的特點(diǎn)就是快。

一個(gè)例子感受一下。

比如求矩陣的前三次冪的和，用NumPy實(shí)現(xiàn)，計(jì)算需要約478毫秒。

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

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

1、NumPy加速器。NumPy的重要性不用多說(shuō)，用Python搞科學(xué)計(jì)算和機(jī)器學(xué)習(xí)，沒(méi)人離得開(kāi)它，但它原生一直不支持GPU等硬件加速。

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

2、XLA。XLA（Accelerated Linear Algebra）就是加速線性代數(shù)，一個(gè)優(yōu)化編譯器。JAX建立在XLA之上，大幅提高了JAX計(jì)算速度的上限。

3、JIT。研究人員可使用XLA將自己的函數(shù)轉(zhuǎn)換為實(shí)時(shí)編譯（JIT）版本，相當(dāng)于通過(guò)向計(jì)算函數(shù)添加一個(gè)簡(jiǎn)單的函數(shù)修飾符，就可以將計(jì)算速度提高幾個(gè)數(shù)量級(jí)。

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

當(dāng)然，JAX也是有一些缺點(diǎn)在身上的。

比如：

1、雖然JAX以加速器著稱，但它并沒(méi)有針對(duì)CPU計(jì)算中的每個(gè)操作進(jìn)行充分優(yōu)化。

2、JAX還太新，沒(méi)有形成像TensorFlow那樣完整的基礎(chǔ)生態(tài)。因此它還沒(méi)有被谷歌以成型產(chǎn)品的形式推出。

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

4、不支持Windows系統(tǒng)，只能在上面的虛擬環(huán)境中運(yùn)行。

5、沒(méi)有數(shù)據(jù)加載器，得借用TensorFlow或PyTorch的。

……

盡管如此，簡(jiǎn)單、靈活又好用的JAX還是率先在DeepMind中流行起來(lái)。2020年誕生的一些深度學(xué)習(xí)庫(kù)Haiku和RLax等都是基于它開(kāi)發(fā)。

這一年，PyTorch原作者之一Adam Paszke，也全職加入了JAX團(tuán)隊(duì)。

目前，JAX的開(kāi)源項(xiàng)目在GitHub上已有18.4k標(biāo)星，比TensorFlow高了不少了。

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

一方面是因?yàn)镴AX的實(shí)力，另一方面主要還是跟TensorFlow自身的很多原因有關(guān)。

為什么谷歌要轉(zhuǎn)投JAX？

誕生于2015年的TensorFlow，曾經(jīng)也風(fēng)靡一時(shí)，推出后很快超過(guò)了Torch、Theano和Caffe等一眾“弄潮兒”，成為最受歡迎的機(jī)器學(xué)習(xí)框架。

然而在2017年，煥然一新的PyTorch“卷土重來(lái)”。

這是Meta基于Torch搭建的機(jī)器學(xué)習(xí)庫(kù)，由于上手簡(jiǎn)單、通俗易懂，很快受到一眾研究者的青睞，甚至有超過(guò)TensorFlow的趨勢(shì)。

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

（從Stack Overflow上的提問(wèn)占比來(lái)看，PyTorch逐年上升，TensorFlow卻一直停滯不前）

在競(jìng)爭(zhēng)之中，TensorFlow的缺點(diǎn)逐漸暴露出來(lái)，API不穩(wěn)定、實(shí)現(xiàn)復(fù)雜、學(xué)習(xí)成本高等問(wèn)題并沒(méi)有隨著更新解決多少，反而結(jié)構(gòu)變得更復(fù)雜了。

相比之下，TensorFlow卻沒(méi)有繼續(xù)發(fā)揮比較能打的“運(yùn)行效率”等優(yōu)勢(shì)。

在學(xué)術(shù)界，PyTorch的使用率正逐漸超過(guò)TensorFlow。

尤其是在各大頂會(huì)如ACL、ICLR中，使用PyTorch實(shí)現(xiàn)的算法框架近幾年已經(jīng)占據(jù)了超過(guò)80%，相比之下TensorFlow的使用率還在不斷下降。

也正是因此，谷歌坐不住了，試圖用JAX奪回對(duì)機(jī)器學(xué)習(xí)框架的“主導(dǎo)權(quán)”。

雖然JAX名義上不是“專為深度學(xué)習(xí)構(gòu)建的通用框架”，然而從發(fā)布之初起，谷歌的資源就一直在向JAX傾斜。

一方面，谷歌大腦和DeepMind逐漸將更多的庫(kù)構(gòu)建在JAX上。

包括谷歌大腦的Trax、Flax、Jax-md，以及DeepMind的神經(jīng)網(wǎng)絡(luò)庫(kù)Haiku和強(qiáng)化學(xué)習(xí)庫(kù)RLax等，都是基于JAX構(gòu)建的。

據(jù)谷歌官方表示：

JAX生態(tài)系統(tǒng)開(kāi)發(fā)中，也會(huì)考慮確保其與現(xiàn)有TensorFlow庫(kù)（如Sonnet和TRFL）的設(shè)計(jì)（盡可能）保持一致。

另一方面，更多的項(xiàng)目也開(kāi)始基于JAX實(shí)現(xiàn)，最近爆火的DALL·E mini項(xiàng)目就是其中一個(gè)。

由于能更好地利用谷歌TPU的優(yōu)勢(shì)，JAX在運(yùn)行性能上比PyTorch要好得多，更多之前搭建在TensorFlow上的工業(yè)界項(xiàng)目也正在轉(zhuǎn)投JAX。

甚至有網(wǎng)友調(diào)侃JAX如今爆火的原因：可能是TensorFlow的使用者實(shí)在無(wú)法忍受這個(gè)框架了。

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

更看好哪一個(gè)框架？

總體來(lái)看，很多人還是很堅(jiān)定地站PyTorch。

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

“JAX雖然很吸引人，但還不夠具備“革命性”的能力促使大家拋棄PyTorch來(lái)使用它?！?/p>

但看好JAX的也并非少數(shù)。

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

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

不過(guò)，總有不怎么care誰(shuí)輸誰(shuí)贏的人，他們的目光放得很長(zhǎng)遠(yuǎn)：

沒(méi)有最好，只有更好。最重要的是更多玩家和好的idea統(tǒng)統(tǒng)都加入進(jìn)來(lái)，讓開(kāi)源和真正優(yōu)秀的創(chuàng)新畫上等號(hào)。

項(xiàng)目地址：
https://github.com/google/jax