you can see each compilation if you use JAX_LOG_COMPILES variable or you use low enough logging level.
Sorry, not to belabor this point.
Would that suggest to you what you did wrong? Or purely show you what you got right? How chatty is this variable?
I used this to see if something is repeatedly compiled. I.e. I have the code that runs in a loop and you immediately see if something is compiled only once, or every time. (and it produces a lot of output) I'm not saying this is the best way to do it though, it just worked for me.