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.cc557 HloSharding GatherIndexSharding(const HloSharding& output_sharding, in GatherIndexSharding()
928 const HloSharding& output_sharding, const HloInstruction& gather, in GatherParallelDataOperandSharding()
1010 const HloSharding& output_sharding, const HloInstruction& hlo) { in GatherDataOperandShardingFromOutput()
1075 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()
479 const HloSharding& output_sharding, in PartitionGather()
Ddot_handler.cc513 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionBaseCase()
1739 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionDotGroupOnBatch()
1971 const HloSharding& output_sharding, in GetNonContractingPartitionGroupedShardingForMatchedOperand()
2006 const HloSharding& other_sharding, const HloSharding& output_sharding, in GetNonContractingPartitionGroupedShardingForOtherOperand()
2077 const Shape& output_base_shape, const HloSharding& output_sharding, in PartitionDotGroupOnNonContracting()
2169 const Shape& output_base_shape, const HloSharding& output_sharding, in GetDotGroupPartitionContractingOutputShardings()
2265 const Shape& output_base_shape, const HloSharding& output_sharding, in PartitionDotGroupOnContracting()
2484 const HloSharding& output_sharding, const SpmdPartitionerOptions& options, in EstimateWindowedEinsumIterationsForNonContractingPartitioning()
2618 const HloSharding& output_sharding, const SpmdPartitionerOptions& options, in PrioritizeContractingDimensionsPartitioning()
2765 const HloSharding& output_sharding, const SpmdPartitionerOptions& options, in LhsIsBestMatchForNonContractingPartitioning()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Dxla_sharding_util.cc327 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