dm-haiku

收录理由

直接用 JAX 搭建神经网络时，一旦涉及定义和复用层，事情往往变得别扭，因为 JAX 的纯函数模型不会自动把参数绑定到使用它们的对象上。Haiku 用一套面向对象的模块接口解决了这个问题，同时仍然产出 JAX 的 jit、grad、pmap 变换所期望的纯 init 和 apply 函数。它出自 Sonnet 原班作者之手，所以任何有 TensorFlow 经验的人都会觉得 API 非常眼熟，而且它曾被用来复现 DeepMind 在视觉、语言、生成模型和强化学习等领域的众多实验。基于这段历史，它作为一层轻量而可靠的工具，适合那些想要模块化组合、但不需要完整框架所附带优化器和检查点机制的科研人员。需要注意的是，Google DeepMind 如今推荐新项目改用 Flax，Haiku 已进入维护模式，因此它主要面向既有 Haiku 代码库，或者作为一份清晰的参考实现来使用。