Searched refs:k_dim (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/core/grappler/costs/ |
| D | op_level_cost_estimator.cc | 1057 double m_dim, n_dim, k_dim, k_dim_b = 0; in CountMatMulOperations() local
1075 k_dim = a_matrix_shape.dim(0).size(); in CountMatMulOperations()
1078 k_dim = a_matrix_shape.dim(1).size(); in CountMatMulOperations()
1088 VLOG(1) << "M, N, K: " << m_dim << "," << n_dim << "," << k_dim; in CountMatMulOperations()
1091 if (k_dim_b != 1 && k_dim != 1 && k_dim_b != k_dim) { in CountMatMulOperations()
1096 k_dim = std::max(k_dim, k_dim_b); in CountMatMulOperations()
1099 ops = m_dim * n_dim * k_dim * 2; in CountMatMulOperations()
1105 mat_mul->k = k_dim; in CountMatMulOperations()
1185 TensorShapeProto_Dim k_dim; in GenerateBatchMatmulContextFromEinsum() local
1189 k_dim.set_size(1); in GenerateBatchMatmulContextFromEinsum()
[all …]
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | nn_ops.cc | 1298 DimensionHandle k_dim; in TopKShapeFn() local
1300 TF_RETURN_IF_ERROR(c->MakeDimForScalarInput(1, &k_dim)); in TopKShapeFn()
1307 k_dim = c->MakeDim(k); in TopKShapeFn()
1311 if (c->ValueKnown(last_dim) && c->ValueKnown(k_dim) && in TopKShapeFn()
1312 c->Value(last_dim) < c->Value(k_dim)) { in TopKShapeFn()
1314 "input must have last dimension >= k = ", c->Value(k_dim), " but is ", in TopKShapeFn()
1321 TF_RETURN_IF_ERROR(c->Concatenate(s, c->Vector(k_dim), &s)); in TopKShapeFn()
|