Hard to find documentation of predefined jax checkpoint policies

[georgedahl]

Description

I tried searching on https://jax.readthedocs.io/en/latest/search.html?q=checkpoint_dots for checkpoint_dots but couldn't find any documentation. The documentation of jax.checkpoint is helpful, but should provide more information on the predefined policies in jax.checkpoint_policies.

What jax/jaxlib version are you using?

No response

Which accelerator(s) are you using?

No response

Additional System Info

No response