今天看代码看到一个新框架,Flax,由Google研究人员开发的
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的概念。