Understanding device memory allocation in Jax

[JINKEHE]

Hello,

I am new to Jax so this might be a silly question.
My question is about memory allocation on GPU when running jitted functions.

First, do I understand correctly that when a jitted function is run, device memory will be allocated at the beginning or during this function for those intermediate variables, such as grads? And then such memory will be released after the end of the function if they are not part of the outputs?

If so, then wouldn't this cause the issue of memory fragmentation on the device if we have multiple jitted functions and each of them runs from time to time such as in RL settings?

In Jax, is there a way then to pin memory for these intermediate variables within a jitted function to avoid memory fragmentation? This avoids memory allocation.

Thank you very much for any help.