Home
last modified time | relevance | path

Searched refs:dim_numbers (Results 1 – 25 of 28) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dconvolution_group_converter.cc201 auto dim_numbers = convolution->convolution_dimension_numbers(); in HandleBatchGroupCount() local
217 int64 input_batch_dimension = dim_numbers.input_batch_dimension(); in HandleBatchGroupCount()
218 const int64 input_feature_dimension = dim_numbers.input_feature_dimension(); in HandleBatchGroupCount()
220 int64 output_batch_dimension = dim_numbers.output_batch_dimension(); in HandleBatchGroupCount()
221 int64 output_feature_dimension = dim_numbers.output_feature_dimension(); in HandleBatchGroupCount()
224 dim_numbers.kernel_input_feature_dimension(); in HandleBatchGroupCount()
226 dim_numbers.kernel_output_feature_dimension(); in HandleBatchGroupCount()
242 for (auto& d : *dim_numbers.mutable_input_spatial_dimensions()) { 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()
[all …]
Dgather_expander.cc111 HloInstruction* index_vector, const GatherDimensionNumbers& dim_numbers, in ExpandIndexVectorIntoOperandSpace() argument
132 int64 index_vector_dim_index = FindIndex(dim_numbers.start_index_map(), i); in ExpandIndexVectorIntoOperandSpace()
133 if (index_vector_dim_index != dim_numbers.start_index_map_size()) { in ExpandIndexVectorIntoOperandSpace()
153 const GatherDimensionNumbers& dim_numbers = gather.gather_dimension_numbers(); in GatherLoopBody() local
161 dim_numbers.index_vector_dim() == in GatherLoopBody()
197 ExpandIndexVectorIntoOperandSpace(index_vector, dim_numbers, in GatherLoopBody()
207 AsInt64Slice(dim_numbers.collapsed_slice_dims()))); in GatherLoopBody()
235 const GatherDimensionNumbers& dim_numbers) { in CreateGatherLoopAccumulatorInitValue() argument
240 if (!absl::c_binary_search(dim_numbers.collapsed_slice_dims(), i)) { in CreateGatherLoopAccumulatorInitValue()
276 const GatherDimensionNumbers& dim_numbers = in GatherLoopTripCount() local
[all …]
Dscatter_expander.cc133 HloInstruction* index_vector, const ScatterDimensionNumbers& dim_numbers, in ExpandIndexVectorIntoOperandSpace() argument
154 FindIndex(dim_numbers.scatter_dims_to_operand_dims(), i); in ExpandIndexVectorIntoOperandSpace()
156 dim_numbers.scatter_dims_to_operand_dims_size()) { in ExpandIndexVectorIntoOperandSpace()
222 const ScatterDimensionNumbers& dim_numbers = in ScatterLoopBody() local
258 ExpandIndexVectorIntoOperandSpace(index_vector, dim_numbers, in ScatterLoopBody()
279 AsInt64Slice(dim_numbers.inserted_window_dims()))); in ScatterLoopBody()
333 const ScatterDimensionNumbers& dim_numbers = in ScatterTripCount() local
337 if (i != dim_numbers.index_vector_dim()) { in ScatterTripCount()
366 const ScatterDimensionNumbers& dim_numbers = in ExpandInstruction() local
389 scatter_indices, dim_numbers.index_vector_dim())); in ExpandInstruction()
[all …]
Dindexed_array_analysis.cc254 const Shape& shape, const GatherDimensionNumbers& dim_numbers, in ComputeArrayForGather() argument
256 if (dim_numbers.index_vector_dim() != indices->shape().dimensions_size()) { in ComputeArrayForGather()
261 CHECK_EQ(dim_numbers.start_index_map_size(), 1); in ComputeArrayForGather()
266 if (dim_numbers.collapsed_slice_dims_size() != 1 || in ComputeArrayForGather()
267 dim_numbers.collapsed_slice_dims(0) != dim_numbers.start_index_map(0)) { in ComputeArrayForGather()
279 if (i != dim_numbers.collapsed_slice_dims(0) && in ComputeArrayForGather()
285 << dim_numbers.collapsed_slice_dims(0); in ComputeArrayForGather()
290 int64 source_dim = dim_numbers.start_index_map(0); in ComputeArrayForGather()
293 if (!absl::c_binary_search(dim_numbers.offset_dims(), i)) { in ComputeArrayForGather()
1033 const Shape& shape, const DotDimensionNumbers& dim_numbers, in ComputeArrayForDotWithIndexedLhs() argument
[all …]
Dbatch_dot_simplification.cc44 const DotDimensionNumbers& dim_numbers = batch_dot->dot_dimension_numbers(); in ElideDegenerateBatchDimensionFromBatchDot() local
52 if (dim_numbers.lhs_contracting_dimensions_size() != 1) { in ElideDegenerateBatchDimensionFromBatchDot()
57 for (int64 batch_dim : dim_numbers.lhs_batch_dimensions()) { in ElideDegenerateBatchDimensionFromBatchDot()
72 DotDimensionNumbers new_dim_numbers = dim_numbers; in ElideDegenerateBatchDimensionFromBatchDot()
76 for (int64 i = 0, e = dim_numbers.lhs_batch_dimensions_size() - in ElideDegenerateBatchDimensionFromBatchDot()
Dshape_inference.cc3213 const GatherDimensionNumbers& dim_numbers) { in ValidateGatherDimensionNumbers() argument
3214 if (!absl::c_is_sorted(dim_numbers.offset_dims())) { in ValidateGatherDimensionNumbers()
3217 StrJoin(dim_numbers.offset_dims(), ", ")); in ValidateGatherDimensionNumbers()
3220 if (absl::c_adjacent_find(dim_numbers.offset_dims()) != in ValidateGatherDimensionNumbers()
3221 dim_numbers.offset_dims().end()) { in ValidateGatherDimensionNumbers()
3224 StrJoin(dim_numbers.offset_dims(), ", ")); in ValidateGatherDimensionNumbers()
3227 const int64 output_offset_dim_count = dim_numbers.offset_dims_size(); in ValidateGatherDimensionNumbers()
3231 for (int i = 0; i < dim_numbers.offset_dims_size(); ++i) { in ValidateGatherDimensionNumbers()
3232 int64 offset_dim = dim_numbers.offset_dims(i); in ValidateGatherDimensionNumbers()
3242 if (dim_numbers.start_index_map_size() != in ValidateGatherDimensionNumbers()
[all …]
Dspace_to_batch_converter.cc80 ConvolutionDimensionNumbers& dim_numbers);
117 HloInstruction* activations, ConvolutionDimensionNumbers& dim_numbers,
162 HloInstruction* activations, ConvolutionDimensionNumbers& dim_numbers,
278 ConvolutionDimensionNumbers dim_numbers = in IsConvSuitableForSpaceToBatch() local
282 if (dim_numbers.input_spatial_dimensions_size() < 1) { in IsConvSuitableForSpaceToBatch()
298 const ConvDetails c = GetConvolutionDetails(convolution, dim_numbers); in IsConvSuitableForSpaceToBatch()
323 int64 activations_batch_dim = dim_numbers.input_batch_dimension(); in IsConvSuitableForSpaceToBatch()
479 HloInstruction* activations, ConvolutionDimensionNumbers& dim_numbers, in BringSpaceNextToBatch() argument
486 ConvolutionDimensionNumbers new_dim_numbers = dim_numbers; in BringSpaceNextToBatch()
503 if (i == dim_numbers.kernel_output_feature_dimension()) { in BringSpaceNextToBatch()
[all …]
Dtriangular_solve_expander.cc77 GatherDimensionNumbers dim_numbers; in DiagonalBlocks() local
79 dim_numbers.add_offset_dims(i); in DiagonalBlocks()
80 dim_numbers.add_start_index_map(i); in DiagonalBlocks()
84 dim_numbers.add_offset_dims(ndims - 1); in DiagonalBlocks()
85 dim_numbers.add_offset_dims(ndims); in DiagonalBlocks()
86 dim_numbers.add_start_index_map(ndims - 2); in DiagonalBlocks()
87 dim_numbers.add_start_index_map(ndims - 1); in DiagonalBlocks()
88 dim_numbers.set_index_vector_dim(1); in DiagonalBlocks()
89 diag_blocks = Gather(a, start_indices, dim_numbers, slice_sizes); in DiagonalBlocks()
Dindexed_array_analysis.h265 const Shape& shape, const GatherDimensionNumbers& dim_numbers,
269 const Shape& shape, const DotDimensionNumbers& dim_numbers,
274 const Shape& shape, const DotDimensionNumbers& dim_numbers,
279 const DotDimensionNumbers& dim_numbers,
Dhlo_cost_analysis_test.cc977 GatherDimensionNumbers dim_numbers; in TEST_F() local
978 dim_numbers.add_offset_dims(1); in TEST_F()
979 dim_numbers.add_collapsed_slice_dims(0); in TEST_F()
980 dim_numbers.add_start_index_map(0); in TEST_F()
981 dim_numbers.set_index_vector_dim(1); in TEST_F()
982 Gather(operand, indices, dim_numbers, {1, 3}); in TEST_F()
1009 ScatterDimensionNumbers dim_numbers; in TEST_F() local
1010 dim_numbers.set_index_vector_dim(1); in TEST_F()
1011 dim_numbers.add_update_window_dims(1); in TEST_F()
1012 dim_numbers.add_inserted_window_dims(0); in TEST_F()
[all …]
Dhlo_evaluator.cc383 const DotDimensionNumbers& dim_numbers, in EvaluateDotOp() argument
393 lhs.shape(), rhs.shape(), dim_numbers, in EvaluateDotOp()
398 dim_numbers, precision_config); in EvaluateDotOp()
1451 const Shape& output_shape, const GatherDimensionNumbers& dim_numbers) { in IterationSpaceForOutputBatchIndices() argument
1458 !absl::c_binary_search(dim_numbers.offset_dims(), i); in IterationSpaceForOutputBatchIndices()
1470 const GatherDimensionNumbers& dim_numbers) { in IterationSpaceForOutputOffsetIndices() argument
1476 absl::c_binary_search(dim_numbers.offset_dims(), i); in IterationSpaceForOutputOffsetIndices()
1478 while (absl::c_binary_search(dim_numbers.collapsed_slice_dims(), in IterationSpaceForOutputOffsetIndices()
1500 const GatherDimensionNumbers* dim_numbers, const Shape& input_shape, in OutputBatchIndexToInputIndex() argument
1502 : dim_numbers_(*dim_numbers), start_indices_(*start_indices) { in OutputBatchIndexToInputIndex()
[all …]
Dhlo_evaluator_typed_visitor.h2082 const Shape& updates_shape, const ScatterDimensionNumbers& dim_numbers) {
2088 !absl::c_binary_search(dim_numbers.update_window_dims(), i);
2101 const Shape& updates_shape, const ScatterDimensionNumbers& dim_numbers) {
2107 absl::c_binary_search(dim_numbers.update_window_dims(), i);
2129 const ScatterDimensionNumbers* dim_numbers, const Shape& input_shape,
2131 : dim_numbers_(*dim_numbers), scatter_indices_(*scatter_indices) {
2259 const ScatterDimensionNumbers& dim_numbers, const Shape& input_shape,
2264 if (absl::c_binary_search(dim_numbers.update_window_dims(), i)) {
2273 if (absl::c_binary_search(dim_numbers.inserted_window_dims(), i)) {
2331 const ScatterDimensionNumbers& dim_numbers =
[all …]
Delemental_ir_emitter.cc1914 const GatherDimensionNumbers& dim_numbers = hlo->gather_dimension_numbers(); in EmitElementalGather() local
1932 if (absl::c_binary_search(dim_numbers.collapsed_slice_dims(), i)) { in EmitElementalGather()
1935 int64 output_window_dim = dim_numbers.offset_dims(operand_index_dim++); in EmitElementalGather()
1945 if (!absl::c_binary_search(dim_numbers.offset_dims(), i)) { in EmitElementalGather()
1954 dim_numbers.index_vector_dim(), in EmitElementalGather()
1971 int64 operand_dim = dim_numbers.start_index_map(dim); in EmitElementalGather()
2000 if (indices_shape.dimensions_size() == dim_numbers.index_vector_dim()) { in EmitElementalGather()
2008 indices_shape.dimensions(dim_numbers.index_vector_dim()); in EmitElementalGather()
2010 gather_index_index_components[dim_numbers.index_vector_dim()] = in EmitElementalGather()
2177 const DotDimensionNumbers& dim_numbers = hlo->dot_dimension_numbers(); in EmitElementalDot() local
[all …]
Dhlo_creation_utils.cc288 const DotDimensionNumbers& dim_numbers, in MakeDotHlo() argument
295 ShapeInference::InferDotOpShape(lhs->shape(), rhs->shape(), dim_numbers, in MakeDotHlo()
298 dot_shape, lhs, rhs, dim_numbers, precision_config)); in MakeDotHlo()
Dhlo_creation_utils.h139 const DotDimensionNumbers& dim_numbers,
Dhlo_evaluator.h127 StatusOr<Literal> EvaluateDotOp(const DotDimensionNumbers& dim_numbers,
/external/tensorflow/tensorflow/compiler/tf2xla/lib/
Dscatter.cc138 xla::ScatterDimensionNumbers dim_numbers; in XlaScatter() local
139 dim_numbers.set_index_vector_dim(indices_are_vectors in XlaScatter()
165 dim_numbers.add_update_window_dims(i); in XlaScatter()
170 dim_numbers.add_inserted_window_dims(i); in XlaScatter()
171 dim_numbers.add_scatter_dims_to_operand_dims(i); in XlaScatter()
193 VLOG(3) << " index_vector_dim: " << dim_numbers.index_vector_dim(); in XlaScatter()
195 << absl::StrJoin(dim_numbers.update_window_dims(), ",") << "]"; in XlaScatter()
197 << absl::StrJoin(dim_numbers.inserted_window_dims(), ",") << "]"; in XlaScatter()
199 << absl::StrJoin(dim_numbers.scatter_dims_to_operand_dims(), ",") in XlaScatter()
203 dim_numbers); in XlaScatter()
/external/tensorflow/tensorflow/compiler/xla/tests/
Ddynamism_inference_test.cc281 GatherDimensionNumbers dim_numbers; in TEST_F() local
282 dim_numbers.add_offset_dims(1); in TEST_F()
283 dim_numbers.add_start_index_map(0); in TEST_F()
284 dim_numbers.set_index_vector_dim(1); in TEST_F()
285 auto gather = Gather(operand1, indices, dim_numbers, {1}); in TEST_F()
301 GatherDimensionNumbers dim_numbers; in TEST_F() local
302 dim_numbers.add_offset_dims(1); in TEST_F()
303 dim_numbers.add_start_index_map(0); in TEST_F()
304 dim_numbers.set_index_vector_dim(1); in TEST_F()
305 auto gather = Gather(data_operand, indices, dim_numbers, {1}); in TEST_F()
[all …]
Dgather_operation_test.cc757 GatherDimensionNumbers dim_numbers; in XLA_TEST_F() local
758 dim_numbers.add_offset_dims(1); in XLA_TEST_F()
759 dim_numbers.add_collapsed_slice_dims(0); in XLA_TEST_F()
760 dim_numbers.add_start_index_map(0); in XLA_TEST_F()
761 dim_numbers.set_index_vector_dim(1); in XLA_TEST_F()
762 Gather(operand, indices, dim_numbers, {1, 3}); in XLA_TEST_F()
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dgather_op.cc122 xla::GatherDimensionNumbers dim_numbers; in XlaGather() local
128 dim_numbers.add_collapsed_slice_dims(i); in XlaGather()
137 dim_numbers.add_offset_dims(i); in XlaGather()
141 dim_numbers.add_offset_dims(i + indices_rank - num_index_dims); in XlaGather()
145 dim_numbers.set_index_vector_dim(indices_are_nd ? (indices_shape.dims() - 1) in XlaGather()
148 dim_numbers.add_start_index_map(i); in XlaGather()
151 *gather_output = xla::Gather(input, indices, dim_numbers, slice_sizes); in XlaGather()
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dmatrix.cc149 xla::GatherDimensionNumbers dim_numbers; in GetMatrixDiagonalViaGather() local
155 dim_numbers.add_collapsed_slice_dims(i); in GetMatrixDiagonalViaGather()
156 dim_numbers.add_start_index_map(i); in GetMatrixDiagonalViaGather()
159 dim_numbers.add_offset_dims(i); in GetMatrixDiagonalViaGather()
165 dim_numbers.set_index_vector_dim(1); in GetMatrixDiagonalViaGather()
167 return Gather(x, start_indices, dim_numbers, slice_sizes, in GetMatrixDiagonalViaGather()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_conv_rewriter.cc72 auto dim_numbers = conv->convolution_dimension_numbers(); in ConvertBatchGroupedToFeatureGroupedConvolution() local
76 int64 input_batch_dimension = dim_numbers.input_batch_dimension(); in ConvertBatchGroupedToFeatureGroupedConvolution()
79 int64 input_feature_dimension = dim_numbers.input_feature_dimension(); in ConvertBatchGroupedToFeatureGroupedConvolution()
117 new_conv->set_convolution_dimension_numbers(dim_numbers); in ConvertBatchGroupedToFeatureGroupedConvolution()
Dir_emission_utils.cc112 const DotDimensionNumbers& dim_numbers = dot.dot_dimension_numbers(); in IsMatrixMultiplication() local
117 dim_numbers.lhs_batch_dimensions_size())) { in IsMatrixMultiplication()
121 CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)), in IsMatrixMultiplication()
122 rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0))); in IsMatrixMultiplication()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Ddot_op_emitter.cc1092 Status ValidateDotDimensionNumbers(const DotDimensionNumbers& dim_numbers) { in ValidateDotDimensionNumbers() argument
1095 TF_RET_CHECK(dim_numbers.lhs_contracting_dimensions_size() == 1); in ValidateDotDimensionNumbers()
1096 std::vector<int64> batch_dim_numbers(dim_numbers.lhs_batch_dimensions_size()); in ValidateDotDimensionNumbers()
1099 absl::c_equal(batch_dim_numbers, dim_numbers.lhs_batch_dimensions())); in ValidateDotDimensionNumbers()
1101 absl::c_equal(batch_dim_numbers, dim_numbers.rhs_batch_dimensions())); in ValidateDotDimensionNumbers()
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/
Dlegalize_to_linalg.cc1201 mhlo::DotDimensionNumbers dim_numbers = op.dot_dimension_numbers(); in matchAndRewrite() local
1203 Extract1DVector(dim_numbers.lhs_batching_dimensions()); in matchAndRewrite()
1205 Extract1DVector(dim_numbers.rhs_batching_dimensions()); in matchAndRewrite()
1207 Extract1DVector(dim_numbers.lhs_contracting_dimensions()); in matchAndRewrite()
1209 Extract1DVector(dim_numbers.rhs_contracting_dimensions()); in matchAndRewrite()

12