Searched refs:dims_a (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | horizontal_input_fusion.cc | 67 auto dims_a = shape_a.dimensions(); in CompareShapeDimsFromLeftToRight() local
69 for (size_t i = 0; i < dims_a.size(); ++i) { in CompareShapeDimsFromLeftToRight()
70 if (dims_a[i] != dims_b[i]) { in CompareShapeDimsFromLeftToRight()
71 return dims_a[i] < dims_b[i]; in CompareShapeDimsFromLeftToRight()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | unroll_batch_matmul.cc | 220 const int dims_a = lhs_shape.size(); in matchAndRewrite() local
222 if (dims_a < 2 || dims_b < 2) { in matchAndRewrite()
243 if (lhs_shape[dims_a - 1] != rhs_shape[dims_b - 2]) { in matchAndRewrite()
248 if (dims_a == 2 && dims_b == 2) { in matchAndRewrite()
293 createMatMulOps(sliced_lhs, sliced_rhs, bcast, lhs_shape[dims_a - 2], in matchAndRewrite()
300 result_shape.push_back(lhs_shape[dims_a - 2]); in matchAndRewrite()
|
| D | batchmatmul_to_einsum.cc | 65 const int dims_a = lhs_shape.size(); in matchAndRewrite() local
67 if (dims_a < 2 || dims_b < 2) { in matchAndRewrite()
|
| /external/tensorflow/tensorflow/lite/toco/graph_transformations/ |
| D | unroll_batch_matmul.cc | 169 const int dims_a = input_array_a.shape().dimensions_count(); in Run() local
171 CHECK_GE(dims_a, 2) << "First input must have rank >= 2"; in Run()
179 CHECK_EQ(input_array_a.shape().dims(dims_a - 1), in Run()
185 if (dims_a == 2 && dims_b == 2) { in Run()
250 result_shape.push_back(input_array_a.shape().dims(dims_a - 2)); in Run()
|
| /external/tensorflow/tensorflow/core/grappler/costs/ |
| D | op_level_cost_estimator_test.cc | 137 OpContext DescribeXlaEinsum(const std::vector<int>& dims_a, in DescribeXlaEinsum() argument
146 if (!dims_a.empty()) in DescribeXlaEinsum()
147 DescribeArbitraryRankInput(dims_a, DT_FLOAT, &op_context.op_info); in DescribeXlaEinsum()
154 OpContext DescribeEinsum(const std::vector<int>& dims_a, in DescribeEinsum() argument
157 OpContext op_context = DescribeXlaEinsum(dims_a, dims_b, equation); in DescribeEinsum()
592 OpContext DescribeBatchMatMul(const std::vector<int>& dims_a, in DescribeBatchMatMul() argument
598 DescribeArbitraryRankInput(dims_a, DT_FLOAT, &op_context.op_info); in DescribeBatchMatMul()
|