Searched refs:padded_hlo (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | fft_handler.cc | 383 auto padded_hlo = PadEachPartitionWithHaloExchange( in HandleFft() local
389 if (padded_hlo.has_value()) { in HandleFft()
390 result = padded_hlo.value(); in HandleFft()
|
| D | spmd_partitioner.cc | 468 auto padded_hlo = in ReshardNoCache() local
472 shard_shape, padded_hlo, in ReshardNoCache()
755 auto padded_hlo = ShapeUtil::Compatible(padded_shape, base_shape_) in ReshardAsWindowedInput() local
761 shard_shape, padded_hlo, offsets_on_padded_shape, in ReshardAsWindowedInput()
1016 auto padded_hlo = PadFromPartialReplicateShape( in ReshardFromPartialReplicateWithDynamicSlice() local
1020 if (!padded_hlo.has_value()) { in ReshardFromPartialReplicateWithDynamicSlice()
1042 shard_shape, padded_hlo.value(), offsets, shard_shape.dimensions())); in ReshardFromPartialReplicateWithDynamicSlice()
1147 auto padded_hlo = PadToShape(hlo_, padded_shape, state_.b); in ReshardWithAllToAll() local
1175 dimensions.push_back(padded_hlo->shape().dimensions(i) / group_size); in ReshardWithAllToAll()
1177 dimensions.push_back(padded_hlo->shape().dimensions(i)); in ReshardWithAllToAll()
[all …]
|
| D | spmd_partitioner_util.cc | 1724 auto padded_hlo = PadBaseShapeBeforeUnevenTiledSharding( in PerGroupSliceFromReplicated() local
1729 shard_shape, padded_hlo, in PerGroupSliceFromReplicated()
|