新智元推荐

来源：AI前线(ID:ai-front)

原文:Reddit策划编辑:Natalie

整理:Vincent编辑:Debra

TensorFlow有了替代品，竟然还是谷歌自己做出来的？这其实是TensorFlow的一个简化库，名为JAX，可以支持部分TensorFlow的功能，但是比TensorFlow更加简洁易用。

什么？TensorFlow有了替代品？什么？竟然还是谷歌自己做出来的？先别慌，从各种意义上来说，这个所谓的“替代品”其实是TensorFlow的一个简化库，名为JAX，结合Autograd和XLA，可以支持部分TensorFlow的功能，但是比TensorFlow更加简洁易用。

虽然还不至于替代TensorFlow，但已经有Reddit网友对JAX寄予厚望，并表示“早就期待能有一个可以直接调用NumpyAPI接口的库了！”，“希望它可以取代TensorFlow！”。

JAX结合了Autograd和XLA，是专为高性能机器学习研究打造的产品。

有了新版本的Autograd，JAX能够自动对Python和NumPy的自带函数求导，支持循环、分支、递归、闭包函数求导，而且可以求三阶导数。它支持自动模式反向求导（也就是反向传播）和正向求导，且二者可以任意组合成任何顺序。

JAX的创新之处在于，它基于XLA在GPU和TPU上编译和运行NumPy程序。默认情况下，编译是在底层进行的，库调用能够及时编译和执行。但是JAX还允许使用单一函数APIjit将自己的Python函数及时编译成经过XLA优化的内核。编译和自动求导可以任意组合，因此可以在不脱离Python环境的情况下实现复杂算法并获得最优性能。

JAX最初由MattJohnson、RoyFrostig、DougalMaclaurin和ChrisLeary发起，他们均任职于谷歌大脑团队。在GitHub的说明文档中，作者明确表示：JAX目前还只是一个研究项目，不是谷歌的官方产品，因此可能会有一些bug。从作者的GitHub简介来看，这应该是谷歌大脑正在尝试的新项目，在同一个GitHub目录下的开源项目还包括8月份在业内引起热议的强化学习框架Dopamine。

以下是JAX的简单使用示例。

GitHub项目传送门：

转载请注明:http://www.aideyishus.com/lkgx/4255.html