Home
last modified time | relevance | path

Searched refs:output_base_shape (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc42 PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape, in PartitionConvolutionWithBatchGroupCount() argument
67 std::vector<int64> rhs_to_lhs_indices(output_base_shape.rank()); in PartitionConvolutionWithBatchGroupCount()
78 std::vector<int64> lhs_to_rhs_indices(output_base_shape.rank()); in PartitionConvolutionWithBatchGroupCount()
127 return PartitionedHlo(sharded_conv, output_base_shape, lhs.state()) in PartitionConvolutionWithBatchGroupCount()
134 PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape, in PartitionConvolutionWithFeatureGroupCount() argument
158 std::vector<int64> rhs_to_lhs_indices(output_base_shape.rank()); in PartitionConvolutionWithFeatureGroupCount()
169 std::vector<int64> lhs_to_rhs_indices(output_base_shape.rank()); in PartitionConvolutionWithFeatureGroupCount()
175 std::vector<int64> lhs_to_output_indices(output_base_shape.rank()); in PartitionConvolutionWithFeatureGroupCount()
219 return PartitionedHlo(sharded_conv, output_base_shape, lhs.state()) in PartitionConvolutionWithFeatureGroupCount()
228 PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() argument
[all …]
Ddot_handler.cc467 PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape, in PartitionBaseCase() argument
493 output_base_shape.rank()); in PartitionBaseCase()
527 return PartitionedHlo(dot, output_base_shape, lhs.state()) in PartitionBaseCase()
603 MakePartitionedShape(output_base_shape, output_sharding); in PartitionBaseCase()
623 LiteralUtil::Zero(output_base_shape.element_type())))); in PartitionBaseCase()
627 LiteralUtil::Zero(output_base_shape.element_type()))); in PartitionBaseCase()
1476 ShapeSizeInBytes(output_base_shape), in PartitionBaseCase()
1498 LiteralUtil::Zero(output_base_shape.element_type()))); in PartitionBaseCase()
1521 MakeBinaryAdd(output_base_shape.element_type(), module)); in PartitionBaseCase()
1523 return PartitionedHlo(ar, output_base_shape, lhs.state()) in PartitionBaseCase()
[all …]
Dconvolution_handler.h30 PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape,