How can I modify the memory exchange mechanism of jax to control the swapping in and out of operators at the bottom level

[Mewo518]

I have written a memory optimization strategy that uses swap technology. It is necessary to control certain nodes in the training graph to be immediately swapped out and swapped in after a period of time. This involves memory management implementation. Which part of the code should be searched for and which part of the content should be modified? I have seen the device_put() function in api.py. Can this function meet my requirements