jax

eadee4f8 - [mosaic] Removed usages of the deprecated `tpu.wait_dma` op

[mosaic] Removed usages of the deprecated `tpu.wait_dma` op The op itself is still present in the dialect and handled in the Mosaic passes. Reverts d2391fcc0e08361d6387935270746967ac1c7f69 PiperOrigin-RevId: 818577739