helx:jax stax的帮助程序库 源码

赫尔克斯 Helx是 / 的帮助程序库。它处于持续发展中,因此请期待变化。有关某些功能,请参见下文。 安装 pip install git+https://github.com/epignatelli/helx 请注意,即使使用helx,它也不依赖于JAX。这使您可以使用为任意加速器构建的JAX版本。 ###贡献该库旨在用作jax和stax的快捷方式的容器。随时提出问题,或请求新功能。 module装饰器 module是stax函数的简单接口。它接收一个stax层构造输出-一个可调用的元组-并返回一个Module对象-一个具有init和apply属性的NamedTuple 。这是一个例子: import jax from jax . experimental . stax import Dense , Relu from helx import Module , module @ mo