openxla / xla

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

[XLA:SPMD] Fix sharding propagation for kGetTupleElement not generating correct sharding shape with subgroup manual. #14304

Open copybara-service[bot] opened 2 days ago

copybara-service[bot] commented 2 days ago

[XLA:SPMD] Fix sharding propagation for kGetTupleElement not generating correct sharding shape with subgroup manual.