jax.example_libraries 模块# JAX 提供了一些用于机器学习的小型实验性库。这些库一方面提供工具,另一方面也作为示例,展示如何使用 JAX 构建此类库。每个库都只有不到 300 行源代码,所以请仔细查看并根据需要进行调整! 注意 每个迷你库都旨在成为一种启发,而不是一种规定。 为了达到这个目的,最好保持它们的示例代码最小化;因此,我们通常不会合并 PR 来添加新功能。相反,请将您优秀的拉取请求和设计想法发送到功能更完整的库,例如 Haiku 或 Flax。 jax.example_libraries.optimizers 模块 jax.example_libraries.stax 模块