jax.ad_checkpoint 模块

checkpoint_name(x, name)

在 jax.checkpoint() 中使用名称标识一个值。