Peak Memory Usage Much Smaller than Peak Heap Usage in JIT?

[Qiustander]

Hi all I am using tensorflow probability with JAX backend. Specifically, I am using particle filter for testing. I used tensorboard to do the profiling for jit-compiled function, and I found that the peak memory usage is much smaller than the peak heap usage, as shown in the figure. I am quite confused, is there anything wrong? Or is there a bug?