Adam Paszke 5c952a1c04 [Mosaic] Swap the order of all-reduce and elementwise ops ...

This lets us only perform an all-reduce once at the end of a reduction, instead of
at every step. This also bundles two small improvements, making layout inference
less strict for `vector.broadcast` and relaxing an assert in elementwise rule.

PiperOrigin-RevId: 552413179

2023-07-31 02:01:57 -07:00

..

dialect/tpu

[Mosaic] Swap the order of all-reduce and elementwise ops

2023-07-31 02:01:57 -07:00

python

[Mosaic] Swap the order of all-reduce and elementwise ops

2023-07-31 02:01:57 -07:00

BUILD

Enable building jaxlib w/ Mosaic

2023-07-26 03:59:30 -07:00