Searched refs:rhs_collapsing_dim (Results 1 – 1 of 1) sorted by relevance
1165 int64 rhs_collapsing_dim = dnums.rhs_contracting_dimensions(0); in HandleDotStrengthReduction() local1166 int64 rhs_kept_dim = kept_dim(rhs_rank, rhs_collapsing_dim, in HandleDotStrengthReduction()1174 std::swap(rhs_collapsing_dim, rhs_kept_dim); in HandleDotStrengthReduction()1225 if (rhs_rank == 2 && rhs->shape().dimensions(rhs_collapsing_dim) == 1) { in HandleDotStrengthReduction()1250 rhs_collapsing_dim), in HandleDotStrengthReduction()1252 rhs_collapsing_dim)))); in HandleDotStrengthReduction()1300 (rhs_collapsing_dim == rhs_rank - 1 && rhs_kept_dim == rhs_rank - 2))) { in HandleDotStrengthReduction()1301 CHECK_EQ(rhs->shape().dimensions(rhs_collapsing_dim), 1); in HandleDotStrengthReduction()1317 broadcast(StripDim(rhs, rhs_collapsing_dim), in HandleDotStrengthReduction()1329 const int64 reduce_dim = rhs_kept_dim < rhs_collapsing_dim in HandleDotStrengthReduction()[all …]