Searched refs:base_shape_ (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 395 element, ShapeUtil::GetTupleElementShape(base_shape_, i), state_) in ReshardNoCache()
402 return PartitionedHlo(tuple, base_shape_, state_); in ReshardNoCache()
451 return PartitionedHlo(copy, base_shape_, state_); in ReshardNoCache()
476 return PartitionedHlo(slice, base_shape_, state_); in ReshardNoCache()
485 if (sharding.IsReplicated() || EvenlyPartitions(base_shape_, sharding)) { in PadWithValue()
501 int64 index_limit = base_shape_.dimensions(dim); in PadWithValue()
517 auto offsets = MakePartitionOffsets(base_shape_, sharding, in PadWithValue()
520 if (base_shape_.dimensions(i) % sharding.tile_assignment().dim(i) == 0 || in PadWithValue()
542 return PartitionedHlo(result, base_shape_, state_); in PadWithValue()
568 auto shard_shape = base_shape_; in ReshardAsWindowedInput()
[all …]
|
| D | spmd_partitioner.h | 294 : hlo_(hlo), base_shape_(base_shape), state_(std::move(state)) { in PartitionedHlo()
324 const Shape& base_shape() const { return base_shape_; } in base_shape()
376 Shape base_shape_; variable
|