Home
last modified time | relevance | path

Searched refs:partitioned_hlo (Results 1 – 2 of 2) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.cc924 auto partitioned_hlo = *this; in ReshardToPartialReplicateWithAllGather() local
927 partitioned_hlo = in ReshardToPartialReplicateWithAllGather()
928 partitioned_hlo.ReshardWithCollectivePermute(temp_sharding); in ReshardToPartialReplicateWithAllGather()
947 partitioned_hlo.hlo_, base_shape_, temp_sharding, target, replicate_dims, in ReshardToPartialReplicateWithAllGather()
948 partitioned_hlo.state().collective_ops_creator, in ReshardToPartialReplicateWithAllGather()
949 partitioned_hlo.state().next_channel_id, in ReshardToPartialReplicateWithAllGather()
950 partitioned_hlo.state().partition_id, partitioned_hlo.state().b); in ReshardToPartialReplicateWithAllGather()
959 partitioned_hlo.state(), sharding_grouped.device_groups, in ReshardToPartialReplicateWithAllGather()
960 partitioned_hlo.state().b); in ReshardToPartialReplicateWithAllGather()
964 auto original_sharding = partitioned_hlo.sharding(); in ReshardToPartialReplicateWithAllGather()
[all …]
Dspmd_partitioner.h462 const PartitionedHlo& partitioned_hlo) { in SetPartitionedHlo() argument
464 partitioned_instructions_.emplace(hlo, partitioned_hlo); in SetPartitionedHlo()