We are observing unexplained out of GPU memory events when trying to train a complex large model (involving conditional execution) and enabling XLA (jit_compile=True flag for tf.function).
Unfortunately, we haven’t been able to reproduce the issue in a reduced shareable form just yet, so I am writing here mostly for feedback.
What we see:

In general GPU utilization for an XLA compiled model goes considerably down compared to non-compiled graph mode or eager execution. This is what we measure in most of our models and in all our small test-cases.
However, in some instances, large models exceed memory capacity when compiled while they can still run in eager mode for the same exact batch size.

Both behaviors seem contradictory, so we are wondering if there are some known corner cases involving XLA that may produce this and we can avoid (we really need the extra train efficiency coming out of XLA).
Thanks!