jax.ad_checkpoint module

checkpoint_name(x, name)

Identifies a value with a name within jax.checkpoint().