Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dgather_scatter_handler.cc169 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()
H A Dconvolution_handler.cc43 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()
H A Ddot_handler.cc759 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 …]
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dhlo_sharding_util_test.cc75 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
H A Dhlo_sharding_util.cc627 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()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
H A Dxla_sharding_util.cc323 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
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dset_default_sharding.cc42 llvm::SmallVector<llvm::StringRef, 4> output_sharding( in SetDefaultSharding() local
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/
H A Ddtensor_meta_ops.cc76 const std::vector<int32> output_sharding = output_layout.num_shards(); in __anon376f01380102() local
/aosp_15_r20/external/pytorch/torch/distributed/tensor/
H A D_op_schema.py457 output_sharding: Optional[OutputSharding] = None variable in OpInfo