Searched refs:dim_nums (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | convolution_dimension_numbers_test.cc | 105 ConvolutionDimensionNumbers dim_nums = in XLA_TEST_F() local
108 int64 old_input_batch_dim = dim_nums.input_batch_dimension(); in XLA_TEST_F()
109 int64 old_output_batch_dim = dim_nums.output_batch_dimension(); in XLA_TEST_F()
110 dim_nums.set_input_batch_dimension(dim_nums.input_feature_dimension()); in XLA_TEST_F()
111 dim_nums.set_output_batch_dimension(dim_nums.output_feature_dimension()); in XLA_TEST_F()
112 dim_nums.set_input_feature_dimension(old_input_batch_dim); in XLA_TEST_F()
113 dim_nums.set_output_feature_dimension(old_output_batch_dim); in XLA_TEST_F()
116 dim_nums.kernel_input_feature_dimension(); in XLA_TEST_F()
117 dim_nums.set_kernel_input_feature_dimension( in XLA_TEST_F()
118 dim_nums.kernel_output_feature_dimension()); in XLA_TEST_F()
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_matchers.cc | 222 const DotDimensionNumbers& dim_nums = instruction->dot_dimension_numbers(); in MatchAndExplain() local
223 if (dim_nums.lhs_contracting_dimensions_size() != 1 || in MatchAndExplain()
224 dim_nums.lhs_contracting_dimensions(0) != lhs_contracting_dim_) { in MatchAndExplain()
227 << absl::StrJoin(dim_nums.lhs_contracting_dimensions(), ",") in MatchAndExplain()
232 if (dim_nums.rhs_contracting_dimensions_size() != 1 || in MatchAndExplain()
233 dim_nums.rhs_contracting_dimensions(0) != rhs_contracting_dim_) { in MatchAndExplain()
236 << absl::StrJoin(dim_nums.rhs_contracting_dimensions(), ",") in MatchAndExplain()
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | dot_op_emitter.cc | 60 DotDimensionNumbers dim_nums; member
69 dim_nums = instr.dot_dimension_numbers(); in DotInfo()
411 const DotDimensionNumbers& dim_nums = dot_info_.dim_nums; in EmitNaiveLlvmIrGemm() local
416 int64 lhs_reduction_dimension = dim_nums.lhs_contracting_dimensions(0); in EmitNaiveLlvmIrGemm()
417 int64 rhs_reduction_dimension = dim_nums.rhs_contracting_dimensions(0); in EmitNaiveLlvmIrGemm()
689 const DotDimensionNumbers& dim_nums = dot_info_.dim_nums; in GetMatMultDims() local
692 /*m=*/lhs_shape.dimensions(1 - dim_nums.lhs_contracting_dimensions(0)), in GetMatMultDims()
693 /*k=*/lhs_shape.dimensions(dim_nums.lhs_contracting_dimensions(0)), in GetMatMultDims()
694 /*n=*/rhs_shape.dimensions(1 - dim_nums.rhs_contracting_dimensions(0)), in GetMatMultDims()
696 /*lhs_non_canonical=*/dim_nums.lhs_contracting_dimensions(0) == 0, in GetMatMultDims()
[all …]
|
| D | dot_op_emitter_internal.h | 41 DotDimensionNumbers dim_nums; member
48 dim_nums = instr.dot_dimension_numbers(); in DotInfo()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | gemm_thunk.cc | 323 DotDimensionNumbers dim_nums = GetDimensionNumbers(*hlo_instruction()); in ExecuteOnStream() local
324 CHECK_EQ(dim_nums.lhs_batch_dimensions_size(), in ExecuteOnStream()
325 dim_nums.rhs_batch_dimensions_size()); in ExecuteOnStream()
326 CHECK_EQ(dim_nums.lhs_batch_dimensions_size() + 2, output_shape_.rank()); in ExecuteOnStream()
328 int64 row_dim = dim_nums.lhs_batch_dimensions_size(); in ExecuteOnStream()
329 int64 col_dim = dim_nums.lhs_batch_dimensions_size() + 1; in ExecuteOnStream()
335 for (int64 batch_dim : dim_nums.lhs_batch_dimensions()) { in ExecuteOnStream()
388 lhs_data, lhs_shape_, dim_nums.lhs_contracting_dimensions(0) == row_dim); in ExecuteOnStream()
390 rhs_data, rhs_shape_, dim_nums.rhs_contracting_dimensions(0) == col_dim); in ExecuteOnStream()
|
| D | gpu_layout_assignment.cc | 194 const DotDimensionNumbers& dim_nums = in AddBackendConstraints() local
196 CHECK_EQ(dim_nums.lhs_batch_dimensions_size(), in AddBackendConstraints()
197 dim_nums.rhs_batch_dimensions_size()); in AddBackendConstraints()
198 CHECK_EQ(dim_nums.lhs_batch_dimensions_size() + 2, in AddBackendConstraints()
200 for (int64 batch_dim : dim_nums.lhs_batch_dimensions()) { in AddBackendConstraints()
|