Searched refs:a_input_shape (Results 1 – 1 of 1) sorted by relevance
290 TensorShapeProto a_input_shape = MaybeGetMinimumShape( in IsEinsumCorrectlyFormed() local297 if (a_input_str.size() != static_cast<size_t>(a_input_shape.dim_size()) || in IsEinsumCorrectlyFormed()1155 TensorShapeProto a_input_shape = MaybeGetMinimumShape( in GenerateBatchMatmulContextFromEinsum() local1199 m_dim.set_size(m_dim.size() * a_input_shape.dim(i_idx).size()); in GenerateBatchMatmulContextFromEinsum()1204 k_dim.set_size(k_dim.size() * a_input_shape.dim(i_idx).size()); in GenerateBatchMatmulContextFromEinsum()1209 *(a_matrix_shape->add_dim()) = a_input_shape.dim(i_idx); in GenerateBatchMatmulContextFromEinsum()1210 *(b_matrix_shape->add_dim()) = a_input_shape.dim(i_idx); in GenerateBatchMatmulContextFromEinsum()1273 TensorShapeProto a_input_shape = MaybeGetMinimumShape( in CountBatchMatMulOperations() local1287 TensorShapeProto* bigger_rank_shape = &a_input_shape; in CountBatchMatMulOperations()1289 if (b_input_shape.dim_size() > a_input_shape.dim_size()) { in CountBatchMatMulOperations()[all …]