| 269 | } |
| 270 | |
| 271 | Status HloSharding::ValidateTuple(const Shape& shape, int64 num_devices) const { |
| 272 | if (!shape.IsTuple()) { |
| 273 | return tensorflow::errors::InvalidArgument( |
| 274 | StrCat("Sharding is tuple-shaped but validation shape is not.")); |
| 275 | } |
| 276 | TF_RETURN_IF_ERROR(CheckLeafCount(shape)); |
| 277 | |
| 278 | // Now we've validated the number of tuple elements, it's safe to request a |
| 279 | // shape tree. |
| 280 | ShapeTree<HloSharding> shape_tree = GetAsShapeTree(shape); |
| 281 | for (const auto& index_to_sharding : shape_tree.leaves()) { |
| 282 | Status status = index_to_sharding.second.ValidateNonTuple( |
| 283 | ShapeUtil::GetSubshape(shape, index_to_sharding.first), num_devices); |
| 284 | if (!status.ok()) { |
| 285 | tensorflow::errors::AppendToMessage( |
| 286 | &status, StrCat("Note: While validating sharding tuple element ", |
| 287 | index_to_sharding.first.ToString(), " which is ", |
| 288 | index_to_sharding.second.ToString())); |
| 289 | return status; |
| 290 | } |
| 291 | } |
| 292 | return Status::OK(); |
| 293 | } |
| 294 | |
| 295 | Status HloSharding::Validate(const Shape& shape, int64 num_devices) const { |
| 296 | Status status = IsTuple() ? ValidateTuple(shape, num_devices) |
nothing calls this directly
no test coverage detected