Searched refs:batch_mat_mul (Results 1 – 3 of 3) sorted by relevance
530 BatchMatMulDimensions* batch_mat_mul, in CountBatchMatMulOperations() argument532 return estimator_.CountBatchMatMulOperations(op_info, batch_mat_mul, in CountBatchMatMulOperations()611 BatchMatMulDimensions batch_mat_mul; in CountBatchMatMulDimProduct() local613 batch_mat_mul.matmul_dims.n = 0; in CountBatchMatMulDimProduct()614 batch_mat_mul.matmul_dims.m = 0; in CountBatchMatMulDimProduct()615 batch_mat_mul.matmul_dims.k = 0; in CountBatchMatMulDimProduct()618 op_info, &batch_mat_mul, found_unknown_shapes); in CountBatchMatMulDimProduct()620 for (auto dim : batch_mat_mul.batch_dims) dimension_product *= dim; in CountBatchMatMulDimProduct()622 dimension_product *= batch_mat_mul.matmul_dims.n; in CountBatchMatMulDimProduct()623 dimension_product *= batch_mat_mul.matmul_dims.m; in CountBatchMatMulDimProduct()[all …]
246 BatchMatMulDimensions* batch_mat_mul,
1240 const OpInfo& op_info, BatchMatMulDimensions* batch_mat_mul, in CountBatchMatMulOperations() argument1302 if (batch_mat_mul != nullptr) { in CountBatchMatMulOperations()1303 batch_mat_mul->batch_dims.push_back(s_dim); in CountBatchMatMulOperations()1342 if (batch_mat_mul != nullptr) { in CountBatchMatMulOperations()1343 batch_mat_mul->matmul_dims.m = (transpose_a.b()) in CountBatchMatMulOperations()1346 batch_mat_mul->matmul_dims.k = (transpose_a.b()) in CountBatchMatMulOperations()1349 batch_mat_mul->matmul_dims.n = (transpose_b.b()) in CountBatchMatMulOperations()