Searched refs:lhs_dim (Results 1 – 6 of 6) sorted by relevance
41 auto broadcast_dim = [](int lhs_dim, int rhs_dim) { in BatchMatMul()42 if (lhs_dim == rhs_dim) return lhs_dim; in BatchMatMul()43 if (lhs_dim == 1) return rhs_dim; in BatchMatMul()45 return lhs_dim; in BatchMatMul()132 auto broadcast_dim = [](int lhs_dim, int rhs_dim) { in BatchMatMul()133 if (lhs_dim == rhs_dim) return lhs_dim; in BatchMatMul()134 if (lhs_dim == 1) return rhs_dim; in BatchMatMul()136 return lhs_dim; in BatchMatMul()286 auto broadcast_dim = [](int lhs_dim, int rhs_dim) { in BatchMatMul()287 if (lhs_dim == rhs_dim) return lhs_dim; in BatchMatMul()[all …]
31 inline int broadcast_dim(int lhs_dim, int rhs_dim) { in broadcast_dim() argument32 if (lhs_dim == rhs_dim) return lhs_dim; in broadcast_dim()33 if (lhs_dim == 1) return rhs_dim; in broadcast_dim()35 return lhs_dim; in broadcast_dim()
107 const int lhs_dim = extended_lhs_shape.Dims(i); in ResizeOutputTensor() local109 int broadcast_dim = lhs_dim; in ResizeOutputTensor()110 if ((lhs_dim != rhs_dim) && (lhs_dim == 1)) { in ResizeOutputTensor()366 const int lhs_dim = extended_lhs_shape.Dims(i); in Prepare() local368 if (lhs_dim != rhs_dim) { in Prepare()369 if (lhs_dim != 1) { in Prepare()
257 for (int64 lhs_dim : dimension_numbers.lhs_batch_dimensions()) { in __anonbbc1c2d10202() local258 output_dims.emplace_back(c->Dim(lhs_shape_handle, lhs_dim)); in __anonbbc1c2d10202()
1277 int64_t lhs_dim = std::get<0>(dim); in GetCompatibleRankedTensorType() local1278 if (lhs_dim == std::get<1>(dim)) { in GetCompatibleRankedTensorType()1279 dims.push_back(lhs_dim); in GetCompatibleRankedTensorType()
709 for (int64 lhs_dim : dimension_numbers.lhs_batch_dimensions()) { in InferDotOpShape() local710 dimensions.push_back(lhs.dimensions(lhs_dim)); in InferDotOpShape()711 is_dynamic.push_back(lhs.is_dynamic_dimension(lhs_dim)); in InferDotOpShape()