jax

2f1645f1 - [mosaic] Added a new tpu.fetch_and_add_sync op

[mosaic] Added a new tpu.fetch_and_add_sync op The op can only target vector subcore SMEM at the moment. Scalar subcore SMEM can be supported in theory, but need more work. PiperOrigin-RevId: 864388263