| 47 | } |
| 48 | |
| 49 | OpSharding Tile1D(const Shape& tile_shape, int64 num_tiles) { |
| 50 | OpSharding result; |
| 51 | result.set_type(OpSharding::OTHER); |
| 52 | |
| 53 | CHECK_EQ(tile_shape.rank(), 1); |
| 54 | std::vector<int64> dimensions(1, num_tiles); |
| 55 | *result.mutable_tile_shape() = tile_shape.ToProto(); |
| 56 | auto& tile_dimension = |
| 57 | (*result.mutable_tile_shape()->mutable_dimensions())[0]; |
| 58 | tile_dimension = CeilOfRatio(static_cast<int64>(tile_dimension), num_tiles); |
| 59 | result.add_tile_assignment_dimensions(num_tiles); |
| 60 | for (int64 i = 0; i < num_tiles; ++i) { |
| 61 | result.add_tile_assignment_devices(i); |
| 62 | } |
| 63 | return result; |
| 64 | } |
| 65 | |
| 66 | OpSharding Tuple(const ShapeTree<OpSharding>& shardings) { |
| 67 | OpSharding result; |
nothing calls this directly
no test coverage detected