jax

af8bdd15 - [XLA:Mosaic] Clean up pass dump and make debug easier.

[XLA:Mosaic] Clean up pass dump and make debug easier. Now each pass dump should only reflect the changes made in its pass. This makes it easier to compare the changes added in a certain pass. PiperOrigin-RevId: 647428192