| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | gather_scatter_handler.cc | 169 const HloSharding& output_sharding, PartitionedHlo& operand, in PartitionGatherIndexPassthroughPartition()
272 const HloSharding& output_sharding, absl::Span<const int64_t> batch_dims, in PartitionGatherPassthroughOperand()
376 const HloSharding& output_sharding, absl::Span<const int64_t> batch_dims, in PartitionGatherTrivialIndexedOperandDimension()
487 const HloSharding& output_sharding, absl::Span<const int64_t> batch_dims, in PartitionGatherIndexParallelDimensions()
640 const HloSharding& output_sharding, absl::Span<const int64_t> batch_dims, in PartitionGather()
776 const Shape& output_shape, const HloSharding& output_sharding, in PartitionScatterPassthroughOperand()
815 const Shape& output_shape, const HloSharding& output_sharding, in PartitionScatterIndexPassthroughPartition()
954 const Shape& output_shape, const HloSharding& output_sharding, in PartitionScatterTrivialIndexedOperandDimension()
1010 const Shape& output_shape, const HloSharding& output_sharding, in PartitionScatter()
|
| D | convolution_handler.cc | 43 const HloSharding& output_sharding, in PartitionConvolutionWithBatchGroupCount()
135 const HloSharding& output_sharding, in PartitionConvolutionWithFeatureGroupCount()
229 const HloSharding& output_sharding, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
518 const HloSharding& output_sharding, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
744 const HloSharding& output_sharding, in PartitionConvolutionTiledOutput()
820 const HloSharding& output_sharding, in PartitionConvolutionBaseCase()
968 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionConvolution()
|
| D | dot_handler.cc | 759 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionBaseCase()
2007 const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping, in PartitionDotGroupOnBatch()
2245 const HloSharding& output_sharding, in GetNonContractingPartitionGroupedShardingForMatchedOperand()
2280 const HloSharding& other_sharding, const HloSharding& output_sharding, in GetNonContractingPartitionGroupedShardingForOtherOperand()
2353 const Shape& output_base_shape, const HloSharding& output_sharding, in PartitionDotGroupOnNonContracting()
2460 const Shape& output_base_shape, const HloSharding& output_sharding, in GetDotGroupPartitionContractingOutputShardings()
2582 const Shape& output_base_shape, const HloSharding& output_sharding, in PartitionDotGroupOnContracting()
2844 const HloSharding& output_sharding, const SpmdPartitionerOptions& options, in EstimateWindowedEinsumIterationsForNonContractingPartitioning()
2989 const HloSharding& output_sharding, const SpmdPartitionerOptions& options, in PrioritizeContractingDimensionsPartitioning()
3217 const HloSharding& output_sharding, const SpmdPartitionerOptions& options, in LhsIsBestMatchForNonContractingPartitioning()
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding_util_test.cc | 75 HloSharding output_sharding = HloSharding::Tile(Array2D<int64_t>({{0}, {1}})); in TEST() local
86 HloSharding output_sharding = in TEST() local
101 HloSharding output_sharding = HloSharding::Tile(tile); in TEST() local
113 HloSharding output_sharding = in TEST() local
139 HloSharding output_sharding = in TEST() local
151 HloSharding output_sharding = in TEST() local
|
| D | hlo_sharding_util.cc | 627 HloSharding GatherIndexSharding(const HloSharding& output_sharding, in GatherIndexSharding()
1012 const HloSharding& output_sharding, const HloInstruction& gather, in GatherParallelDataOperandSharding()
1078 const HloSharding& output_sharding, const HloInstruction& hlo) { in GatherDataOperandShardingFromOutput()
|
| /external/tensorflow/tensorflow/dtensor/mlir/ |
| D | set_default_sharding.cc | 42 llvm::SmallVector<llvm::StringRef, 4> output_sharding( in SetDefaultSharding() local
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | xla_sharding_util.cc | 323 const auto& output_sharding = output_sharding_and_index.value(); in ParseAndValidateOutputSharding() local
472 const xla::OpSharding& output_sharding, in ValidateAndGetTiledExecuteOutputShape()
521 const auto& output_sharding = output_sharding_config[output_index]; in GetOutputTypesForLogicalDeviceComputation() local
564 const auto& output_sharding = output_sharding_config[output_index]; in RemapOutputsFromLogicalDevices() local
|
| /external/tensorflow/tensorflow/dtensor/cc/ |
| D | dtensor_meta_ops.cc | 76 const std::vector<int32> output_sharding = output_layout.num_shards(); in __anon82f11ae70102() local
|