jax

1b23d21c - [Mosaic GPU] Reland back skipping multimem address initialization.

[Mosaic GPU] Reland back skipping multimem address initialization. multimem_addr call is now skipped for all modules which doesn't use any mulimem operations. Reverts af08a3be01b58cbc1acea0005165933b8e8fa882 PiperOrigin-RevId: 965196260