jax

ab81c0f6 - [Mosaic GPU] Use the newly introduced scale copy API to simplify and strengthen the block scaled test

[Mosaic GPU] Use the newly introduced scale copy API to simplify and strengthen the block scaled test The test is both simpler, and tests scales that have completely random patterns. PiperOrigin-RevId: 779158520