Home
last modified time | relevance | path

Searched refs:input_batch_dimension (Results 1 – 25 of 43) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dconvolution_group_converter.cc217 int64 input_batch_dimension = dim_numbers.input_batch_dimension(); in HandleBatchGroupCount() local
229 activation->shape().dimensions(input_batch_dimension); in HandleBatchGroupCount()
238 input_sizes[input_batch_dimension] /= batch_group_count; in HandleBatchGroupCount()
239 input_sizes.insert(input_sizes.begin() + input_batch_dimension, in HandleBatchGroupCount()
243 if (d > input_batch_dimension) { in HandleBatchGroupCount()
247 dim_numbers.add_input_spatial_dimensions(input_batch_dimension); in HandleBatchGroupCount()
248 dim_numbers.set_input_batch_dimension(input_batch_dimension + 1); in HandleBatchGroupCount()
249 if (input_feature_dimension > input_batch_dimension) { in HandleBatchGroupCount()
596 if (dim_numbers.input_batch_dimension() > input_feature_dimension) { in HandleConvolution()
597 dim_numbers.set_input_batch_dimension(dim_numbers.input_batch_dimension() + in HandleConvolution()
Dtranspose_folding_test.cc388 new_conv->convolution_dimension_numbers().input_batch_dimension()); in TEST_F()
390 dnums.input_batch_dimension(), in TEST_F()
456 new_conv->convolution_dimension_numbers().input_batch_dimension()); in TEST_F()
458 dnums.input_batch_dimension(), in TEST_F()
Dconvolution_4d_expander.cc128 removed_input_dimensions, new_dim_nums.input_batch_dimension())); in ExpandInstruction()
Dspace_to_batch_converter.cc323 int64 activations_batch_dim = dim_numbers.input_batch_dimension(); in IsConvSuitableForSpaceToBatch()
535 if (is_backprop && i == dim_numbers.input_batch_dimension()) { in BringSpaceNextToBatch()
772 if (consumer->convolution_dimension_numbers().input_batch_dimension() != in CanPropagate()
1862 DimLookUp(permute_dims, original_conv_dims.input_batch_dimension()); in PropagateOnConv()
1876 const int64 old_batch_dim = original_conv_dims.input_batch_dimension(); in PropagateOnConv()
2187 const int64 prev_batch_dim = original_conv_dims.input_batch_dimension(); in PropagateOnBackpropFilterConv()
2283 DimLookUp(permute_dims, original_conv_dims.input_batch_dimension()); in PropagateOnBackpropFilterConv()
2663 int64 activations_batch_dim = dim_numbers.input_batch_dimension(); in PerformSpaceToBatchOnConvolution()
2771 << " batch dim " << new_dim_numbers.input_batch_dimension(); in PerformSpaceToBatchOnConvolution()
Dtranspose_folding.cc171 transpose_dimensions[dnums.input_batch_dimension()]); in FoldTransposeIntoConvolution()
Ddot_as_convolution_util.cc58 {conv_dims.input_batch_dimension(), -1, in ParseConvolutionDimsInfo()
Ddynamic_dimension_inference.cc569 if (dimension == dimension_numbers.input_batch_dimension()) { in HandleConvolution()
697 if (dimension == dimension_numbers.input_batch_dimension()) { in HandleDynamicConvolutionForward()
Dshape_inference.cc1683 input_dnums[0] = dnums.input_batch_dimension(); in InferConvolveShape()
1739 const int64 input_batch = lhs.dimensions(dnums.input_batch_dimension()); in InferConvolveShape()
1831 if (i == dnums.input_batch_dimension()) { in InferConvolveShape()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_conv_rewriter.cc76 int64 input_batch_dimension = dim_numbers.input_batch_dimension(); in ConvertBatchGroupedToFeatureGroupedConvolution() local
88 reshape_dims[input_batch_dimension] = in ConvertBatchGroupedToFeatureGroupedConvolution()
89 reshape_dims[input_batch_dimension] / num_groups; in ConvertBatchGroupedToFeatureGroupedConvolution()
90 reshape_dims.insert(reshape_dims.begin() + input_batch_dimension, num_groups); in ConvertBatchGroupedToFeatureGroupedConvolution()
98 transpose_dims.erase(transpose_dims.begin() + input_batch_dimension); in ConvertBatchGroupedToFeatureGroupedConvolution()
100 input_batch_dimension); in ConvertBatchGroupedToFeatureGroupedConvolution()
181 auto input_batch_dim = conv_dnums.input_batch_dimension(); in MatchBackwardFilter()
518 dnums.set_output_batch_dimension(conv_dnums.input_batch_dimension()); in MatchBackwardInput()
Dstream_executor_util.cc55 input_layout.push_back(dnums.input_batch_dimension()); in StreamExecutorConvLayoutsToXlaLayouts()
62 input_layout.push_back(dnums.input_batch_dimension()); in StreamExecutorConvLayoutsToXlaLayouts()
Dgpu_conv_runner.cc379 .set_count(input_shape.dimensions(dnums.input_batch_dimension())); in GetGpuConvConfig()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc58 lhs.base_shape().dimensions(dnums.input_batch_dimension()); in PartitionConvolutionWithBatchGroupCount()
69 dnums.input_batch_dimension(); in PartitionConvolutionWithBatchGroupCount()
85 lhs_to_output_indices[dnums.input_batch_dimension()] = in PartitionConvolutionWithBatchGroupCount()
102 (ShardCountAtDim(lhs.sharding(), dnums.input_batch_dimension()) == in PartitionConvolutionWithBatchGroupCount()
162 dnums.input_batch_dimension(); in PartitionConvolutionWithFeatureGroupCount()
178 lhs_to_output_indices[dnums.input_batch_dimension()] = in PartitionConvolutionWithFeatureGroupCount()
242 dnums.input_batch_dimension(); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
262 return lhs_sharding.tile_assignment().dim(dnums.input_batch_dimension()) != in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
293 (lhs.sharding().tile_assignment().dim(dnums.input_batch_dimension()) > in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
532 dnums.input_batch_dimension(); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
[all …]
Ddot_handler.cc232 int64 input_batch_dimension = dnums.input_batch_dimension(); in GenNewConvDNums() local
238 if (lhs_concat_dim <= input_batch_dimension) { in GenNewConvDNums()
239 input_batch_dimension++; in GenNewConvDNums()
307 new_dnums.set_input_batch_dimension(input_batch_dimension); in GenNewConvDNums()
2263 dnums.input_batch_dimension(); in ConvertDimsMappingWithFeatureGroupCount()
2279 new_dims_mapping.batch_dims.back().lhs = dnums.input_batch_dimension(); in ConvertDimsMappingWithBatchGroupCount()
/external/tensorflow/tensorflow/compiler/mlir/hlo/tests/
Dlhlo_gpu_ops.mlir37 { dimension_numbers = {input_batch_dimension = 0 : i64,
72 dimension_numbers = {input_batch_dimension = 0 : i64,
102 dimension_numbers = {input_batch_dimension = 0 : i64,
134 dimension_numbers = {input_batch_dimension = 0 : i64,
165 dimension_numbers = {input_batch_dimension = 0 : i64,
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/translate/
Dlayouts_and_names.mlir11 input_batch_dimension = 0 : i64,
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dir_emission_utils.cc104 return dnums.input_batch_dimension() == 0 && in PotentiallyImplementedAsEigenConvolution()
Dconv_canonicalization.cc41 auto input_batch_dim = dnums.input_batch_dimension(); in Run()
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_client.py539 self.input_batch_dimension = 0
583 dimension_numbers.input_batch_dimension = 0
596 dimension_numbers.input_batch_dimension = lhs_spec.index('N')
/external/tensorflow/tensorflow/compiler/xla/tests/
Dconvolution_dimension_numbers_test.cc108 int64 old_input_batch_dim = dim_nums.input_batch_dimension(); in XLA_TEST_F()
/external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/
Dhlo_ops_base_structs.td42 StructFieldAttr<"input_batch_dimension",I64Attr>,
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dattribute_exporter.cc31 input.input_batch_dimension().getValue().getSExtValue()); in ConvertConvDimensionNumbers()
Dattribute_importer.cc116 builder->getI64IntegerAttr(dnums.input_batch_dimension()), in ConvertConvDimensionNumbers()
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/
Dhlo_text_to_lhlo_no_opt.hlotxt284 // CHECK-SAME: input_batch_dimension = 0 : i64
321 // CHECK-SAME: input_batch_dimension = 0 : i64
360 // CHECK-SAME: input_batch_dimension = 0 : i64
/external/tensorflow/tensorflow/compiler/mlir/xla/experimental/conv_emitter/
Dconv_emitter.cc536 GetShapeInfo(conv->operand(0)->shape(), dim_nums.input_batch_dimension(), in EmitConvolutionForwardAsMlir()
/external/tensorflow/tensorflow/compiler/tests/
Dxla_ops_test.py138 dnums.input_batch_dimension = 0

12