Register-level FLOPS counting in JAX programs ?

[mitkotak]

Hey I am on the hunt for a FLOPS counter that can count up the ADD/MUL happening at the register level since JAX's FLOP counter can be unreliable. I was wondering whether there's any way to do this using nsys-jax.

If yes, then is there a way to set up nsys-jax up outside a docker container i.e are there a list of instructions that I can send to my sysadmin to help them set it up ?

Thanks !

[mitkotak]

Just wanted to share an update. What I have been doing so far is running ncu and then counting up the FLOPS using instructions similar to here but as you can imagine, its extremely tedious to run.
Are there any alternatives out there that would be useful particularly for ML pipelines ? I am also open to some form of XLA/HLO level tracer.

Thanks for your input !