| 85 | } |
| 86 | |
| 87 | string HloSharding::ToString() const { |
| 88 | if (IsTuple()) { |
| 89 | std::vector<string> parts; |
| 90 | parts.reserve(tuple_elements_.size()); |
| 91 | for (const HloSharding& element : tuple_elements_) { |
| 92 | parts.push_back(element.ToString()); |
| 93 | } |
| 94 | return StrCat("{", absl::StrJoin(parts, ", "), "}"); |
| 95 | } |
| 96 | |
| 97 | if (replicated_) { |
| 98 | return "{replicated}"; |
| 99 | } |
| 100 | if (maximal_) { |
| 101 | return StrCat( |
| 102 | "{maximal device=", static_cast<int64>(*tile_assignment_.begin()), "}"); |
| 103 | } |
| 104 | return StrCat("{devices=[", StrJoin(tile_assignment_.dimensions(), ","), "]", |
| 105 | StrJoin(tile_assignment_, ","), "}"); |
| 106 | } |
| 107 | |
| 108 | bool HloSharding::UsesDevice(int64 device) const { |
| 109 | if (IsTuple()) { |
no test coverage detected