JAX

数组计算与程序变换库 · 自动微分、图变换、编译与批处理

Python

实现线索

自动微分重计算

在微分变换中选择保存的中间值与需要重算的部分，控制内存和计算开销。

查看 jax/_src/ad_checkpoint.pyefe43f1bd5