static */
| 1285 | } |
| 1286 | |
| 1287 | /* static */ absl::optional<Shape> ShapeUtil::AlignLayouts( |
| 1288 | const Shape& input_shape, const Shape& output_shape) { |
| 1289 | CHECK(input_shape.IsArray()); |
| 1290 | CHECK(output_shape.IsArray()); |
| 1291 | // Removing trivial dimensions from the shape simplifies the alignment |
| 1292 | // algorithm since ones can go in any position. |
| 1293 | if (HasDegenerateDimensions(input_shape) || |
| 1294 | HasDegenerateDimensions(output_shape)) { |
| 1295 | auto simple_output_shape = |
| 1296 | AlignLayouts(DropDegenerateDimensions(input_shape), |
| 1297 | DropDegenerateDimensions(output_shape)); |
| 1298 | if (!simple_output_shape) { |
| 1299 | return absl::nullopt; |
| 1300 | } |
| 1301 | |
| 1302 | std::vector<int64> layout = |
| 1303 | SpanToVector(simple_output_shape->layout().minor_to_major()); |
| 1304 | // For each one sized dimension in the output, increment the dimension |
| 1305 | // numbers in layout that are more minor than the one. |
| 1306 | absl::InlinedVector<int64, 8> dim_map; |
| 1307 | dim_map.reserve(simple_output_shape->rank()); |
| 1308 | for (int64 i = 0; i < output_shape.rank(); ++i) { |
| 1309 | if (output_shape.dimensions(i) != 1) { |
| 1310 | dim_map.push_back(i); |
| 1311 | } |
| 1312 | } |
| 1313 | for (int64& d : layout) { |
| 1314 | d = dim_map[d]; |
| 1315 | } |
| 1316 | |
| 1317 | // Add the ones in descending order to the layout. Descending layouts tend |
| 1318 | // to reduce the number of copies inserted in layout assignment. |
| 1319 | for (int64 i = output_shape.rank() - 1; i >= 0; --i) { |
| 1320 | if (output_shape.dimensions(i) == 1) { |
| 1321 | layout.push_back(i); |
| 1322 | } |
| 1323 | } |
| 1324 | Shape output_shape_with_layout = output_shape; |
| 1325 | *output_shape_with_layout.mutable_layout() = Layout{layout}; |
| 1326 | return output_shape_with_layout; |
| 1327 | } |
| 1328 | |
| 1329 | int64 input_rank = input_shape.rank(); |
| 1330 | int64 output_rank = output_shape.rank(); |
| 1331 | |
| 1332 | // First, calculate an alignment of the dimensions. A consecutive sequence of |
| 1333 | // input dimensions and output dimensions belong to the same alignment part if |
| 1334 | // the products of their dimension bounds are the same. In the easiest case, |
| 1335 | // an alignment part consists of one input dimension and one output dimension |
| 1336 | // which both have the same dimension bound. An alignment part specifies which |
| 1337 | // dimensions need to be kept together in a physical layout if we want a |
| 1338 | // reshape to be a bitcast. The order of the alignment parts is defined by the |
| 1339 | // physical layout of the input shape, so when we construct the layout for the |
| 1340 | // output shape we just process the alignment parts in this order, and then |
| 1341 | // layout the dimensions belonging to each part in descending (major to minor) |
| 1342 | // order. |
| 1343 | |
| 1344 | // Stores the input and output dimension numbers where each alignment part |
nothing calls this directly
no test coverage detected