| 26 | namespace { |
| 27 | |
| 28 | Status PopulateInfeedLayoutVector(const xla::Shape& shape, |
| 29 | std::vector<int>* layouts) { |
| 30 | if (shape.IsTuple()) { |
| 31 | int64 tuple_elements = xla::ShapeUtil::TupleElementCount(shape); |
| 32 | for (int64 i = 0; i < tuple_elements; ++i) { |
| 33 | const xla::Shape& subshape = |
| 34 | xla::ShapeUtil::GetTupleElementShape(shape, i); |
| 35 | TF_RETURN_IF_ERROR(PopulateInfeedLayoutVector(subshape, layouts)); |
| 36 | } |
| 37 | } else if (xla::LayoutUtil::HasLayout(shape)) { |
| 38 | for (auto dim : xla::LayoutUtil::MinorToMajor(shape)) { |
| 39 | layouts->push_back(dim); |
| 40 | } |
| 41 | } else { |
| 42 | layouts->insert(layouts->end(), shape.rank(), -1); |
| 43 | } |
| 44 | return Status::OK(); |
| 45 | } |
| 46 | |
| 47 | // Populate the output layout unless the minor_to_major array contains all -1 |
| 48 | // value, in which case the layout is considered missing and the API returns |