Searched refs:hlo_ (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.h | 294 : hlo_(hlo), base_shape_(base_shape), state_(std::move(state)) { in PartitionedHlo()
299 if (hlo_->shape().IsTuple() && !hlo_->sharding().IsTuple()) { in PartitionedHlo()
300 hlo_->set_sharding( in PartitionedHlo()
301 hlo_->sharding().GetTupleSharding(hlo_->shape()).ValueOrDie()); in PartitionedHlo()
318 HloInstruction* hlo() const { return hlo_; } in hlo()
321 const HloSharding& sharding() const { return hlo_->sharding(); } in sharding()
373 HloInstruction* hlo_; variable
|
| D | spmd_partitioner.cc | 351 hlo_->shape().IsArray() && target.NumTiles() < sharding().NumTiles(); in Reshard()
370 VLOG(2) << "Resharding " << hlo_->ToString() << " from " in ReshardNoCache()
371 << hlo_->sharding().ToString() << " to " << target.ToString(); in ReshardNoCache()
372 const Shape& shape = hlo_->shape(); in ReshardNoCache()
449 HloInstruction::CreateUnary(hlo_->shape(), HloOpcode::kCopy, hlo_)); in ReshardNoCache()
461 hlo_, state_.partition_id, target_grouped.device_groups, group_dims, in ReshardNoCache()
469 PadBaseShapeBeforeUnevenTiledSharding(hlo_, target, state_.b); in ReshardNoCache()
482 const HloSharding& sharding = hlo_->sharding(); in PadWithValue()
483 const Shape& shape = hlo_->shape(); in PadWithValue()
540 shape, HloOpcode::kSelect, mask, hlo_, broadcast_pad_value)); in PadWithValue()
[all …]
|