How does jax.jit decide when to use gradient checkpointing/recomputation?

[zhixuan-lin]

I read here that jax.jit does gradient checkpointing automatically. How does it decide when to use it for different operations? I'm noticing strange memory usage changes in my program with certain code changes (e.g., broadcasting a small array to a larger array dramatically reduces memory usage). I suspect that this is because it causes jax.jit to make different decisions regarding gradient checkpointing, but I'm not sure.