Is it possible to do gradient checkpointing with Jax?

[nyngwang]

I'm trying to migrate from tensorflow2 to Jax because of some problems, and in a hurry to implement something called gradient-checkpointing. I'm reading this: https://jax.readthedocs.io/en/latest/_autosummary/jax.checkpoint.html?#jax-checkpoint, but I'm not sure whether it can solve my problems or not:

1. Is checkpoint of Jax the same idea as the recompute_grad of tensorflow?: tensorflow has tf.keras to define layers in class. And after all the layers are defined I just wrap those adjacent layers I want to recompute into tf.keras.Sequential to make it work with tf.recompute_grad.
2. How is GPU-memory profiling done in Jax, if I need to show numbers to people that the $O(\sqrt n)$ memory usage can be achieved if those correctly select some "checkpoints"? It seems that tensorflow has its own internal allocator to do some smart strategy to save even more memory so I never know the real usage of GPU memory.