Flax深度学习框架

今天看代码看到一个新框架，Flax，由Google研究人员开发的

Flax是基于JAX上构建的
文档链接
github

JAX

JAX是一个函数转换的集合，例如实时编译和自动微分，它是用一个API在XLA上实现的瘦包装器，API本质上是NumPy和SciPy的替代品。事实上，开始使用JAX的一种方法是将其视为一个加速器支持的NumPy。

import jax.numpy as np

# Will be seamlessly executed on an accelerator such as GPU/TPU.
x, w, b = np.ones((3, 1000, 1000))
y = np.dot(w, x) + b

目前，JAX作者在开发新特性时似乎坚持自己的核心能力。当然，合理的方法也是其主要缺点之一：缺乏内置的神经网络组件，除了证明Stax的概念。