Searched refs:rhs_dim (Results 1 – 4 of 4) 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()44 TFLITE_DCHECK_EQ(rhs_dim, 1); 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()135 TFLITE_DCHECK_EQ(rhs_dim, 1); 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 …]
415 const int64 rhs_dim = rhs_dimension_number(dim_name); in Einsum() local418 if (x_shape.dimensions(i) == y_shape.dimensions(rhs_dim)) { in Einsum()420 dnums.add_rhs_batch_dimensions(rhs_dim); in Einsum()422 rhs_outer_dims.push_back(rhs_dim); in Einsum()426 rhs_delete_dims.push_back(rhs_dim); in Einsum()429 if (x_shape.dimensions(i) == y_shape.dimensions(rhs_dim)) { in Einsum()431 dnums.add_rhs_contracting_dimensions(rhs_dim); in Einsum()433 rhs_outer_dims.push_back(rhs_dim); in Einsum()437 rhs_delete_dims.push_back(rhs_dim); in Einsum()
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()34 TFLITE_DCHECK_EQ(rhs_dim, 1); in broadcast_dim()
108 const int rhs_dim = extended_rhs_shape.Dims(i); in ResizeOutputTensor() local110 if ((lhs_dim != rhs_dim) && (lhs_dim == 1)) { in ResizeOutputTensor()111 broadcast_dim = rhs_dim; in ResizeOutputTensor()367 const int rhs_dim = extended_rhs_shape.Dims(i); in Prepare() local368 if (lhs_dim != rhs_dim) { in Prepare()370 TF_LITE_ENSURE_EQ(context, rhs_dim, 1); in Prepare()