谷歌抛弃 TensorFlow，押宝 JAX，其中都有哪些值得关注的信息？

2022/06/21

大V博文

谷歌抛弃 TensorFlow，押宝 JAX，其中都有哪些值得关注的信息？

作者：金雪锋链接：https://www.zhihu.com/question/537823429/answer/2537694584 现在说谷歌抛弃TensorFlow，押宝JAX，估计还早，有点标题党。从使用者的角度看，当前除了谷歌外，其他公司使用JAX的感觉并不多，从Papers with code网站的统计看，论文的占比还比较少。但是JAX这个技术方向是很有意思的，这里代表了几个趋势：

1、从业务上看，JAX的驱动力是AI+科学计算，和传统的CV/NLP相比，算是一个新赛道，从框架角度看，核心的能力是NN计算和数值计算的融合。

金雪锋：AI框架的演进趋势和MindSpore的构想

2、从接口看，兼容numpy、scipy等python原生的数据科学的接口，在此基础上扩展分布式（pmap/pjit）、向量化（vmap）、grad（各种高级微分）、JIT（硬件加速），既能兼容现有的Python生态，又在此基础上增加可微分能力，同时通过各种机制实现扩展性和性能加速。

3、从编程风格上，全面采用函数式的风格，主要体现在无副作用、lamda闭包等；特别是闭包的特性可以使得各种功能可以灵活组合，比如：pmap（vmap（func））-pmap和vmap组合。另外，需要强调的是微分机制，这个与Pytorch的差异就更大一点，Jax的微分机制是函数式的自动微分机制，而Pytorch的是以tensor为中心的机制，我个人感觉函数式的其实更加符合算法的直观，下面有一个以前的简单分析：

金雪锋：聊一聊AI框架前端

4、实际上函数式这一套，pytorch也在跟随，但是个人觉得，与pytorch原来的设计理念有冲突。

金雪锋：Pytorch 1.11看上去很美，但是.....

不过，JAX目前看，并不算特别成熟，主要表现在，从传统的深度学习的角度看，nn的表达完全使用纯函数并不方便，所以后来jax也发展了很多nn封装的库，比如haiku等；另外，jax在分布式方面做的还不够成熟。总的来说，个人感觉，谷歌在tensorflow发展不顺利的情况下，通过jax来重新定义api标准、在新赛道上（AI+科学计算）构建全新的能力是一种好的尝试。