Home
last modified time | relevance | path

Searched refs:TransposeSharding (Results 1 – 7 of 7) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc97 hlo_sharding_util::TransposeSharding(lhs.sharding(), rhs_to_lhs_indices); in PartitionConvolutionWithBatchGroupCount()
99 hlo_sharding_util::TransposeSharding(rhs.sharding(), lhs_to_rhs_indices); in PartitionConvolutionWithBatchGroupCount()
119 auto aligned_output_sharding = hlo_sharding_util::TransposeSharding( in PartitionConvolutionWithBatchGroupCount()
188 hlo_sharding_util::TransposeSharding(lhs.sharding(), rhs_to_lhs_indices); in PartitionConvolutionWithFeatureGroupCount()
190 hlo_sharding_util::TransposeSharding(rhs.sharding(), lhs_to_rhs_indices); in PartitionConvolutionWithFeatureGroupCount()
212 auto aligned_output_sharding = hlo_sharding_util::TransposeSharding( in PartitionConvolutionWithFeatureGroupCount()
254 hlo_sharding_util::TransposeSharding(lhs.sharding(), rhs_to_lhs_indices); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
256 hlo_sharding_util::TransposeSharding(rhs.sharding(), lhs_to_rhs_indices); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
563 hlo_sharding_util::TransposeSharding(lhs.sharding(), rhs_to_lhs_indices), in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
565 auto aligned_lhs_sharding = hlo_sharding_util::TransposeSharding( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
[all …]
Dspmd_partitioner.cc1129 temp_target_tile = hlo_sharding_util::TransposeSharding( in ReshardWithAllToAll()
1957 hlo_sharding_util::TransposeSharding(sharding, inverse_dimensions); in HandleTranspose()
Dspmd_partitioner_util.cc380 auto transpose_sharding = hlo_sharding_util::TransposeSharding( in PartialReplicateReshardCompatibleSharding()
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_sharding_util_test.cc25 EXPECT_EQ(TransposeSharding(HloSharding::Replicate(), {0, 1, 2}), in TEST()
33 EXPECT_EQ(TransposeSharding(input, {3, 0, 1, 2}), output); in TEST()
Dhlo_sharding_util.h81 HloSharding TransposeSharding(const HloSharding& sharding,
Dhlo_sharding_util.cc268 HloSharding TransposeSharding(const HloSharding& sharding, in TransposeSharding() function
1167 auto transposed = TransposeSharding(sharding, dim_permutation); in PartiallyReplicateTiledShardingOnDims()
1243 auto tgt_sharding = hlo_sharding_util::TransposeSharding(source, perm); in TransposeShardingWithCollapsedDims()
Dsharding_propagation.cc530 return hlo_sharding_util::TransposeSharding(lhs->sharding(), in InferConvolutionShardingFromOperands()
784 HloSharding sharding = hlo_sharding_util::TransposeSharding( in InferShardingFromOperands()
1268 return hlo_sharding_util::TransposeSharding(user.sharding(), in GetShardingFromUser()