AI研究的提速器! DeepMind力荐的JAX到底有多强大？( 二 ) ：AI研究的提速器!DeepMind力荐

文章图片
详情请参看下方链接：
https://github.com/deepmind/dm-haiku
2.Optax:基于梯度的优化方法是机器学习的基础。 Optax提供了包含梯度转换与合成运算符(例如链)的工具库，这使得仅仅一行代码就可以实现许多标准的优化器(例如RMSProp或Adam) 。 Optax天然支持用户在定制优化器中对基本部分进行重组，并且还提供了许多用于随机梯度估计和二阶优化的工具集。

文章图片
详情请参看下方链接：
https://github.com/deepmind/optax
3.RLax：许多成功的项目都位于深度学习与强化学习的交叉学科前沿。 RLax库为构建RL主体提供了有效的基础模块。 RLax中的组件涵盖了丰富的算法和思想，包括：TD学习法、策略梯度、参与者评定法(actorcritics)、MAP、近端策略优化、非线性价值转换以及通用价值函数和多种学习探索方法。

文章图片
尽管RLax提供了一些介绍性的构建代理的例子，但它的主要目标并不是用作构建和部署完整RL系统的框架。另一个工具包Acme是基于RLax组件构建的框架，如果需要框架性的支持可以使用Acem来完成。
详情请参看下方链接：
Rlax：https://github.com/deepmind/rlax
【AI研究的提速器! DeepMind力荐的JAX到底有多强大？】Acem：https://deepmind.com/research/publications/Acme
4.Chex：测试对于验证软件的可靠性至关重要，对于验证代码的正确性也同样重要。代码的正确性将提高研究结论的可信度。为了对代码进行有效测试，测试工具集Chex为工具库作者提供了验证通用构件是否正确且健壮的有效手段，终端用户也可以使用它们来检查实验代码是否可靠。

文章图片
Chex提供了各种实用的工具，包括可识别JAX的单元测试、JAX数据类型的属性声明校验、模拟和伪造多种设备的测试环境。 Chex已被用于DeepMind的JAX生态系统以及多个外部项目(例如Coax和MineRL) 。
详情请参看下方链接：
https://github.com/deepmind/chex
5.Jraph：图神经网络(GNN)是近年来飞速发展的研究领域，在交通流量预测和物流模拟等众多领域具有非常强大的应用潜力。 Jraph(与“长颈鹿”一个发音)是一个轻量级的库，支持在JAX中构建和使用GNN 。

文章图片
Jraph提供了用于图数据结构的标准化数据结构，以及一系列用于处理图的工具库和构建可扩展GNNs的神经网络库。此外还包括下面一些特性：基于GraphTuples的批处理可以有效利用硬件加速器的性能；通过填充、掩膜以及在输入分区上定义损失来实现对于可变形图卷积的JIT支持。与Optax及其他JAX库一样， Jraph对用户选择其他神经网络库没有的任何限制。
详情请参看下方链接：
https://github.com/deepmind/jraph
JAX生态系统一直在持续演进，越来越多的研究与开发者开始使用JAX生态中的工具库来加速研究与开发，相信你也能从中收获到更强大的机器学习性能！
更多细节请参看下面的链接：
https://deepmind.com/blog/article/using-jax-to-accelerate-our-research
https://github.com/google/jax#neural-network-libraries