点上方蓝字计算机视觉联盟获取更多干货

在右上方 ··· 设为星标 ★，与你不见不散

编辑：Sophia
计算机视觉联盟  报道  | 公众号 CVLianMeng

转载于 ：Github, https://github.com/google/jax

【人工智能资源（书籍+视频）全网收集，附下载！】

推荐文章【点击下面可直接跳转】：

如何快速下载不同网站文档！

AI博士笔记系列推荐：

博士笔记 | 周志华《机器学习》手推笔记“神经网络”

JAX是Autograd和 XLA的结合，用于高性能机器学习研究。

通过其更新版本的Autograd，JAX可以自动区分本机Python和NumPy函数。它可以通过循环，分支，递归和闭包来区分，并且可以采用派生类的派生类。它支持反向模式微分（也称为反向传播）grad以及正向模式微分，并且两者可以任意组合为任意顺序。

新功能是JAX使用 XLA 在GPU和TPU上编译和运行您的NumPy程序。默认情况下，编译是在后台进行的，而库调用将得到及时的编译和执行。但是，JAX还允许您使用单功能API即时将自己的Python函数编译为XLA优化的内核 jit。编译和自动微分可以任意组合，因此您可以表达复杂的算法并获得最佳性能，而无需离开Python。

深入研究，您会发现JAX实际上是可组合函数转换的可扩展系统 。这两个 grad和jit 是这种转变的实例。另一个是vmap 用于自动矢量化的功能，将来还会有更多功能。

import jax.numpy as np
from jax import grad, jit, vmapdef predict(params, inputs):for W, b in params:outputs = np.dot(inputs, W) + binputs = np.tanh(outputs)return outputsdef logprob_fun(params, inputs, targets):preds = predict(params, inputs)return np.sum((preds - targets)**2)grad_fun = jit(grad(logprob_fun))  # compiled gradient evaluation function
perex_grads = jit(vmap(grad_fun, in_axes=(None, 0, 0)))  # fast per-example grads

JAX是用纯Python编写的，但是它依赖于XLA，它需要作为jaxlib软件包进行编译和安装。按照以下说明使用来安装二进制软件包pip，或从源代码构建JAX。

我们支持jaxlib在Linux（Ubuntu 16.04或更高版本）和macOS（10.12或更高版本）平台（而不是Windows）上安装或构建。我们目前不支持Windows，但欢迎您提供帮助（请参阅＃438）。一些用户报告说，jaxlib使用Windows Subsytem for Linux从源代码构建纯CPU成功。

END

声明：本文来源于网络

如有侵权，联系删除

联盟学术交流群

扫码添加联盟小编，可与相关学者研究人员共同交流学习：目前开设有人工智能、机器学习、计算机视觉、自动驾驶（含SLAM）、Python、求职面经、综合交流群扫描添加CV联盟微信拉你进群，备注：CV联盟  

最新热文荐读

GitHub | 计算机视觉最全资料集锦

Github | 标星1W+清华大学计算机系课程攻略！

Github | 吴恩达新书《Machine Learning Yearning》

收藏 | 2020年AI、CV、NLP顶会最全时间表！

收藏 | 博士大佬总结的Pycharm 常用快捷键思维导图！

收藏 | 深度学习专项课程精炼图笔记！

笔记 | 手把手教你使用PyTorch从零实现YOLOv3

笔记 | 如何深入理解计算机视觉？（附思维导图）

笔记 | 深度学习综述思维导图（可下载）

笔记 | 深度神经网络综述思维导图（可下载）

点个在看支持一下吧