Home
last modified time | relevance | path

Searched defs:output_sharding (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_sharding_util_test.cc60 HloSharding output_sharding = HloSharding::Tile(Array2D<int64>({{0}, {1}})); in TEST() local
71 HloSharding output_sharding = in TEST() local
86 HloSharding output_sharding = HloSharding::Tile(tile); in TEST() local
98 HloSharding output_sharding = in TEST() local
124 HloSharding output_sharding = in TEST() local
136 HloSharding output_sharding = in TEST() local
Dhlo_sharding_util.cc537 HloSharding GatherIndexSharding(const HloSharding& output_sharding, in GatherIndexSharding()
895 const HloSharding& output_sharding, const HloInstruction& gather, in GatherParallelDataOperandSharding()
962 const HloSharding& output_sharding, const HloInstruction& hlo) { in GatherDataOperandShardingFromOutput()
1032 const HloSharding& output_sharding, const HloInstruction& hlo) { in ScatterUpdateShardingFromOutput()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc43 const HloSharding& output_sharding, in PartitionConvolutionWithBatchGroupCount()
135 const HloSharding& output_sharding, in PartitionConvolutionWithFeatureGroupCount()
229 const HloSharding& output_sharding, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
516 const HloSharding& output_sharding, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
742 const HloSharding& output_sharding, in PartitionConvolutionTiledOutput()
818 const HloSharding& output_sharding, in PartitionConvolutionBaseCase()
966 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionConvolution()
Dgather_scatter_handler.cc181 const HloSharding& output_sharding, absl::Span<const int64> batch_dims, in ParititonPassthroughOperand()
216 const HloSharding& output_sharding, absl::Span<const int64> batch_dims, in ParititonTrivialIndexedOperandDimension()
316 const HloSharding& output_sharding, absl::Span<const int64> batch_dims, in PartitionIndexParallelDimensions()
463 const HloSharding& output_sharding, in PartitionGather()
Ddot_handler.cc468 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionBaseCase()
1650 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionDotGroupOnBatch()
1882 const HloSharding& output_sharding, in GetNonContractingPartitionGroupedShardingForMatchedOperand()
1916 const HloSharding& other_sharding, const HloSharding& output_sharding, in GetNonContractingPartitionGroupedShardingForOtherOperand()
1986 const Shape& output_base_shape, const HloSharding& output_sharding, in PartitionDotGroupOnNonContracting()
2081 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionDotGroupOnContracting()
2290 const HloSharding& output_sharding, const SpmdPartitionerOptions& options, in LhsIsBestMatchForNonContractingPartitioning()
2433 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionDot()
2800 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionDot()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Dxla_sharding_util.cc329 const auto& output_sharding = output_sharding_and_index.value(); in ParseAndValidateOutputSharding() local
461 const xla::OpSharding& output_sharding, in ValidateAndGetTiledExecuteOutputShape()
508 const auto& output_sharding = output_sharding_config[output_index]; in GetOutputTypesForLogicalDeviceComputation() local
545 const auto& output_sharding = output_sharding_config[output_index]; in RemapOutputsFromLogicalDevices() local