jax

fc52bcfa - [Mosaic GPU] Emit a constant multicast mask for (2, 1, 1) collective commit arrivals

[Mosaic GPU] Emit a constant multicast mask for (2, 1, 1) collective commit arrivals Also keep the collective matmul test output in float32 instead of downcasting before comparison against the float32 reference. Signed-off-by: Gabriel Castro <gcastro@nvidia.com>