openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators
Apache License 2.0
2.55k stars 394 forks source link

[XLA] Add support for dot(dot(a, b), c) to dot(a, dot(b, c)) reordering where dimensions of a and b are batched in the inner dot, and the corresponding dimensions are batched with dimensions of c in the outer dot. #14876

Closed copybara-service[bot] closed 3 weeks ago

copybara-service[bot] commented 1 month ago

[XLA] Add support for dot(dot(a, b), c) to dot(a, dot(b, c)) reordering where dimensions of a and b are batched in the inner dot, and the corresponding dimensions are batched with dimensions of c in the outer dot.