search

openxla / xla

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

Fix a bug where two replicated `xla::ifrt::HloSharding`s with different numbers of devices were considered having the same partitioning #14868

Closed copybara-service[bot] closed 3 months ago

copybara-service[bot] commented 3 months ago

Fix a bug where two replicated xla::ifrt::HloShardings with different numbers of devices were considered having the same partitioning

xla::HloSharding does not encode the number of devices for replicated sharding, so we must explicitly compare the device count. This new semantics is consistent with other implementations, e.g., ConcreteEvenSharding also compares the number of devices in addition to the global/shard shapes.