| 128 | } |
| 129 | |
| 130 | Status GetShapeWithLayout( |
| 131 | const xla::Shape& input_shape, absl::Span<const int64> minor_to_major, |
| 132 | const std::function<xla::Layout(const xla::Shape&)>& layout_func, |
| 133 | xla::Shape* output_shape) { |
| 134 | if (input_shape.IsTuple()) { |
| 135 | int64 tuple_elements = xla::ShapeUtil::TupleElementCount(input_shape); |
| 136 | std::vector<xla::Shape> shapes; |
| 137 | shapes.reserve(tuple_elements); |
| 138 | size_t position = 0; |
| 139 | for (int64 i = 0; i < tuple_elements; ++i) { |
| 140 | const xla::Shape& shape = |
| 141 | xla::ShapeUtil::GetTupleElementShape(input_shape, i); |
| 142 | if (shape.IsTuple()) { |
| 143 | return errors::InvalidArgument( |
| 144 | "Nested tuples not supported: ", |
| 145 | xla::ShapeUtil::HumanString(input_shape)); |
| 146 | } |
| 147 | int64 rank = shape.rank(); |
| 148 | if (position + rank > minor_to_major.size()) { |
| 149 | return errors::InvalidArgument( |
| 150 | "Not enough layout attribute elements: position=", position, |
| 151 | " rank=", rank, " elements=", minor_to_major.size()); |
| 152 | } |
| 153 | shapes.push_back(shape); |
| 154 | TF_RETURN_IF_ERROR(AssignLayout( |
| 155 | absl::Span<const int64>(minor_to_major).subspan(position, rank), |
| 156 | layout_func, &shapes.back())); |
| 157 | position += rank; |
| 158 | |
| 159 | VLOG(4) << "Shape[" << i |
| 160 | << "] = " << xla::ShapeUtil::HumanStringWithLayout(shapes.back()); |
| 161 | } |
| 162 | if (position != minor_to_major.size()) { |
| 163 | return errors::InvalidArgument( |
| 164 | "Too many elements passed in the layout attribute: position=", |
| 165 | position, " size=", minor_to_major.size()); |
| 166 | } |
| 167 | *output_shape = xla::ShapeUtil::MakeTupleShape(shapes); |
| 168 | } else { |
| 169 | int64 rank = input_shape.rank(); |
| 170 | if (rank != minor_to_major.size()) { |
| 171 | return errors::InvalidArgument( |
| 172 | "Wrong number of layout attribute elements: rank=", rank, |
| 173 | " elements=", minor_to_major.size()); |
| 174 | } |
| 175 | *output_shape = input_shape; |
| 176 | TF_RETURN_IF_ERROR(AssignLayout(minor_to_major, layout_func, output_shape)); |
| 177 | |
| 178 | VLOG(4) << "Shape[] = " |
| 179 | << xla::ShapeUtil::HumanStringWithLayout(*output_shape); |
| 180 | } |
| 181 | return Status::OK(); |
| 182 | } |
| 183 | |
| 184 | } // namespace tensorflow |
no test coverage detected