Searched refs:tile_assignment_ (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding.cc | 102 "{maximal device=", static_cast<int64>(*tile_assignment_.begin()), "}"); in ToString()
104 return StrCat("{devices=[", StrJoin(tile_assignment_.dimensions(), ","), "]", in ToString()
105 StrJoin(tile_assignment_, ","), "}"); in ToString()
114 const auto& devices = tile_assignment_; in UsesDevice()
145 tile_assignment_.Each([&](absl::Span<const int64> index, int64 d) { in TileIndexForDevice()
158 return *tile_assignment_.begin(); in DeviceForTileIndex()
160 return tile_assignment_(index); in DeviceForTileIndex()
171 CHECK_EQ(shape.dimensions_size(), tile_assignment_.num_dimensions()); in TileOffsetForDevice()
176 index[i] * CeilOfRatio(shape_dim, tile_assignment_.dim(i)), shape_dim); in TileOffsetForDevice()
190 CHECK_EQ(shape.dimensions_size(), tile_assignment_.num_dimensions()); in TileLimitForDevice()
[all …]
|
| D | hlo_sharding.h | 190 tile_assignment_ == other.tile_assignment_ &&
205 const Array<int64>& tile_assignment() const { return tile_assignment_; } in tile_assignment()
223 tile_assignment_({0}) {} in HloSharding()
234 tile_assignment_({1}, device_id) {} in HloSharding()
239 tile_assignment_(tile_assignment) {} in HloSharding()
244 tile_assignment_({0}), in HloSharding()
276 Array<int64> tile_assignment_; variable
|