Searched refs:dims_b (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | unroll_batch_matmul.cc | 221 const int dims_b = rhs_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()
294 rhs_shape[dims_b - 1], element_type, loc, rewriter); in matchAndRewrite()
301 result_shape.push_back(rhs_shape[dims_b - 1]); in matchAndRewrite()
|
| D | batchmatmul_to_einsum.cc | 66 const int dims_b = rhs_shape.size(); in matchAndRewrite() local
67 if (dims_a < 2 || dims_b < 2) { in matchAndRewrite()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | horizontal_input_fusion.cc | 68 auto dims_b = shape_b.dimensions(); in CompareShapeDimsFromLeftToRight() local
70 if (dims_a[i] != dims_b[i]) { in CompareShapeDimsFromLeftToRight()
71 return dims_a[i] < dims_b[i]; in CompareShapeDimsFromLeftToRight()
|
| /external/tensorflow/tensorflow/lite/toco/graph_transformations/ |
| D | unroll_batch_matmul.cc | 170 const int dims_b = input_array_b.shape().dimensions_count(); in Run() local
172 CHECK_GE(dims_b, 2) << "Second input must have rank >= 2"; in Run()
180 input_array_b.shape().dims(dims_b - 2)) in Run()
185 if (dims_a == 2 && dims_b == 2) { in Run()
251 result_shape.push_back(input_array_b.shape().dims(dims_b - 1)); in Run()
|
| /external/tensorflow/tensorflow/core/grappler/costs/ |
| D | op_level_cost_estimator_test.cc | 122 const std::vector<int>& dims_b, in DescribeSparseTensorDenseMatMul() argument
131 DescribeArbitraryRankInput(dims_b, DT_FLOAT, &op_context.op_info); in DescribeSparseTensorDenseMatMul()
138 const std::vector<int>& dims_b, in DescribeXlaEinsum() argument
148 if (!dims_b.empty()) in DescribeXlaEinsum()
149 DescribeArbitraryRankInput(dims_b, DT_FLOAT, &op_context.op_info); in DescribeXlaEinsum()
155 const std::vector<int>& dims_b, in DescribeEinsum() argument
157 OpContext op_context = DescribeXlaEinsum(dims_a, dims_b, equation); in DescribeEinsum()
593 const std::vector<int>& dims_b) { in DescribeBatchMatMul() argument
599 DescribeArbitraryRankInput(dims_b, DT_FLOAT, &op_context.op_info); in DescribeBatchMatMul()
|