mirror of
https://github.com/ROCm/jax.git
synced 2025-04-24 18:36:07 +00:00

- Enable it for minor or second-minor implicit dims for the non-no-op case. - Don't allow output offsets for broadcasted dimensions to be non-replicated. Make sure to assign them as replicated in infer-vector-layout for all cases. - Don't fail when both tiled dimensions are logically broadcasted but only one of them requires actual broadcasting (before, it would hit the unimplemented sublane + lane broadcast case). PiperOrigin-RevId: 637772134