Home
last modified time | relevance | path

Searched refs:conv_dnums (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_conv_rewriter.cc179 const ConvolutionDimensionNumbers& conv_dnums = in MatchBackwardFilter() local
181 auto input_batch_dim = conv_dnums.input_batch_dimension(); in MatchBackwardFilter()
182 auto input_feature_dim = conv_dnums.input_feature_dimension(); in MatchBackwardFilter()
183 auto input_spatial_dims = conv_dnums.input_spatial_dimensions(); in MatchBackwardFilter()
184 auto kernel_input_feature_dim = conv_dnums.kernel_input_feature_dimension(); in MatchBackwardFilter()
185 auto kernel_output_feature_dim = conv_dnums.kernel_output_feature_dimension(); in MatchBackwardFilter()
186 auto kernel_spatial_dims = conv_dnums.kernel_spatial_dimensions(); in MatchBackwardFilter()
187 auto output_batch_dim = conv_dnums.output_batch_dimension(); in MatchBackwardFilter()
188 auto output_feature_dim = conv_dnums.output_feature_dimension(); in MatchBackwardFilter()
189 auto output_spatial_dims = conv_dnums.output_spatial_dimensions(); in MatchBackwardFilter()
[all …]
Dgpu_conv_rewriter_test.cc276 ConvolutionDimensionNumbers conv_dnums; in TEST_F() local
277 conv_dnums.set_input_batch_dimension(0); in TEST_F()
278 conv_dnums.set_output_batch_dimension(0); in TEST_F()
279 conv_dnums.set_input_feature_dimension(1); in TEST_F()
280 conv_dnums.set_output_feature_dimension(1); in TEST_F()
281 conv_dnums.add_input_spatial_dimensions(2); in TEST_F()
282 conv_dnums.add_output_spatial_dimensions(2); in TEST_F()
283 conv_dnums.add_input_spatial_dimensions(3); in TEST_F()
284 conv_dnums.add_output_spatial_dimensions(3); in TEST_F()
285 conv_dnums.set_kernel_input_feature_dimension(0); in TEST_F()
[all …]
Dgpu_conv_padding_legalization.cc49 Window* conv_window, const ConvolutionDimensionNumbers& conv_dnums, in MaybePaddedAndSlicedInput() argument
64 for (size_t i = 0; i < conv_dnums.input_spatial_dimensions().size(); ++i) { in MaybePaddedAndSlicedInput()
65 int64 dim = conv_dnums.input_spatial_dimensions(i); in MaybePaddedAndSlicedInput()
98 for (size_t i = 0; i < conv_dnums.input_spatial_dimensions().size(); ++i) { in MaybePaddedAndSlicedInput()
99 int64 dim = conv_dnums.input_spatial_dimensions(i); in MaybePaddedAndSlicedInput()
124 const ConvolutionDimensionNumbers& conv_dnums, in MaybePaddedKernel() argument
135 for (size_t i = 0; i < conv_dnums.kernel_spatial_dimensions().size(); ++i) { in MaybePaddedKernel()
136 int64 dim = conv_dnums.kernel_spatial_dimensions(i); in MaybePaddedKernel()
/external/tensorflow/tensorflow/compiler/xla/service/
Ddot_as_convolution_util.cc113 const auto& conv_dnums = conv.convolution_dimension_numbers(); in CreateShardedConvForDotGeneralConvolution() local
118 conv_dnums.input_spatial_dimensions(dim.spatial_dim))); in CreateShardedConvForDotGeneralConvolution()
128 conv_dnums.input_spatial_dimensions(dim.spatial_dim))); in CreateShardedConvForDotGeneralConvolution()
136 conv_dnums.kernel_spatial_dimensions(dim.spatial_dim))); in CreateShardedConvForDotGeneralConvolution()
145 /*batch_group_count=*/conv.batch_group_count(), window, conv_dnums, in CreateShardedConvForDotGeneralConvolution()
151 /*batch_group_count=*/conv.batch_group_count(), window, conv_dnums, in CreateShardedConvForDotGeneralConvolution()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc899 const auto& conv_dnums = conv.convolution_dimension_numbers(); in CreateShardedConvConvolution() local
904 conv_dnums.input_spatial_dimensions(dim.spatial_dim))); in CreateShardedConvConvolution()
914 conv_dnums.input_spatial_dimensions(dim.spatial_dim))); in CreateShardedConvConvolution()
922 conv_dnums.kernel_spatial_dimensions(dim.spatial_dim))); in CreateShardedConvConvolution()
938 conv_dnums.input_feature_dimension()) / in CreateShardedConvConvolution()
940 conv_dnums.kernel_input_feature_dimension()); in CreateShardedConvConvolution()
946 sharded_lhs_hlo->shape().dimensions(conv_dnums.input_batch_dimension()); in CreateShardedConvConvolution()
953 feature_group_count, batch_group_count, window, conv_dnums, in CreateShardedConvConvolution()
958 batch_group_count, window, conv_dnums, conv.precision_config()); in CreateShardedConvConvolution()
Ddot_handler.cc143 const ConvolutionDimensionNumbers& conv_dnums = in GenNewWindow() local
146 for (int64 i = 0; i < conv_dnums.input_spatial_dimensions_size(); ++i) { in GenNewWindow()
147 if (conv_dnums.input_spatial_dimensions(i) == lhs_concat_dim) { in GenNewWindow()
166 for (int64 i = 0; i < conv_dnums.kernel_spatial_dimensions_size(); ++i) { in GenNewWindow()
167 if (conv_dnums.kernel_spatial_dimensions(i) == rhs_concat_dim && in GenNewWindow()