JAX 介绍
JAX 是一个 TensorFlow 的简化库，它结合了 Autograd 和 XLA，专门用于高性能机器学习研究。
凭借 Autograd，JAX 可以求导循环、分支、递归和闭包函数，并且它可以进行三阶求导。通过
grad，它支持自动模式反向求导（反向传播）和正向求导，且二者可以任何顺序任意组合。
得力于 XLA，可以在 GPU 和 TPU 上编译和运行 NumPy 程序。默认情况下，编译发生在底层，库调用实时编译和执行。但是 JAX
还允许使用单一函数 API jit 将 Python 函数及时编译为 XLA 优化的内核。编译和自动求导可以任意组合，因此可以在 Python
环境下实现复杂的算法并获得最大的性能。
import jax.numpy as np from jax import grad, jit, vmap from functools import partial def predict(params, inputs): for W, b in params: outputs = np.dot(inputs, W) + b inputs = np.tanh(outputs) return outputs def 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 实际上是一个可扩展的可组合函数转换系统，grad 和 jit 都是这种转换的实例。
JAX 官网
版权声明：本文内容由互联网用户自发贡献，该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容， 请发送邮件至 [email protected] 举报，一经查实，本站将立刻删除。
