Searched refs:GetSubSharding (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding.h | 218 HloSharding GetSubSharding(const Shape& shape, const ShapeIndex& index) const;
|
| D | sharding_propagation.cc | 626 HloSharding new_sharding = operand->sharding().GetSubSharding( in InferShardingFromOperands()
1238 auto sub_sharding = user.sharding().GetSubSharding( in GetShardingFromUser()
1272 auto sub_sharding = user.sharding().GetSubSharding( in GetShardingFromUser()
1315 ? user.sharding().GetSubSharding( in GetShardingFromUser()
1344 return user_sharding = user_sharding.GetSubSharding( in GetShardingFromUser()
|
| D | hlo_sharding_metadata.cc | 176 .GetSubSharding(user.shape(), {user.operand_index(&instruction)}) in GetShardingTreeFromUser()
|
| D | hlo_sharding.cc | 625 HloSharding HloSharding::GetSubSharding(const Shape& shape, in GetSubSharding() function in xla::HloSharding
|
| D | layout_assignment.cc | 916 ? instruction.sharding().GetSubSharding(instruction.shape(), index) in SetupCopiedInstruction()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 392 element->set_sharding(sharding().GetSubSharding(shape, {i})); in ReshardNoCache()
396 .Reshard(target.GetSubSharding(shape, {i})) in ReshardNoCache()
1748 sharding = hlo->sharding().GetSubSharding(hlo->shape(), {0}); in HandleSort()
1750 if (sharding != hlo->sharding().GetSubSharding(hlo->shape(), {i})) { in HandleSort()
2481 tuple.sharding().GetSubSharding(tuple.base_shape(), {hlo->tuple_index()}); in HandleGetTupleElement()
2506 auto sharding = hlo->sharding().GetSubSharding(hlo->shape(), {0}); in HandleInfeed()
2707 per_input_sharding = hlo->sharding().GetSubSharding(hlo->shape(), {0}); in HandleReduce()
3372 .Reshard(hlo->sharding().GetSubSharding(hlo->shape(), {i})) in HandleTuple()
|
| D | spmd_partitioner_util.cc | 131 sharding.GetSubSharding(shape, {i}))) { in EvenlyPartitions()
154 sharding.GetSubSharding(shape, {i}))); in MakePartitionedShape()
174 sharding.GetSubSharding(shape, {i}), partition_id)); in MakeNonPaddedShapeForGivenPartition()
|