在JAX中使用jax.clear_caches()遇错误及缓存排查需求

我使用JAX在GPU上训练神经网络，多轮训练后出现内存不足问题。通过在每轮训练后调用jax.profiler.save_device_memory_profile分析，发现每轮内存都会增加固定的量。由于我对输入大小不同的复杂函数使用了vmap和pmap，怀疑函数每轮都在被重新编译，进而导致内存累积。

我的训练流程大致如下：

for i in range(number_of_epochs):
   self._trainer.train(1)

为验证是否是缓存函数累积导致内存泄漏，我尝试在每轮训练后调用jax.clear_cashes()，修改后的代码如下：

for i in range(number_of_epochs):
   self._trainer.train(1)
   jax.clear_cashes()

但执行后第一轮训练结束就出现错误：

malloc_consolidate(): invalid chunk size

我尚未掌握jax.clear_cashes()的正确用法，希望能得到相关帮助。

补充需求 #

可能仅凭现有信息难以直接给出解决方案，但如果能查看每轮训练后缓存函数的数量，对我排查问题会有很大帮助，请问该如何实现？

内容的提问来源于stack exchange，提问作者Anton B