Searched refs:dim_set (Results 1 – 2 of 2) sorted by relevance
544 const absl::flat_hash_set<int64> dim_set(dims.begin(), dims.end()); in Reduce4DTo1D() local545 CHECK_EQ(dim_set.size(), 3); in Reduce4DTo1D()546 for (int64 a0 = 0; a0 == 0 || (!dim_set.contains(0) && a0 < array.n1()); in Reduce4DTo1D()548 for (int64 a1 = 0; a1 == 0 || (!dim_set.contains(1) && a1 < array.n2()); in Reduce4DTo1D()550 for (int64 a2 = 0; a2 == 0 || (!dim_set.contains(2) && a2 < array.n3()); in Reduce4DTo1D()552 for (int64 a3 = 0; a3 == 0 || (!dim_set.contains(3) && a3 < array.n4()); in Reduce4DTo1D()556 i0 == 0 || (dim_set.contains(0) && i0 < array.n1()); ++i0) { in Reduce4DTo1D()558 i1 == 0 || (dim_set.contains(1) && i1 < array.n2()); ++i1) { in Reduce4DTo1D()560 i2 == 0 || (dim_set.contains(2) && i2 < array.n3()); ++i2) { in Reduce4DTo1D()562 i3 == 0 || (dim_set.contains(3) && i3 < array.n4()); in Reduce4DTo1D()
629 absl::flat_hash_set<int64> dim_set; in ValidateDotDimensionNumbers() local630 auto is_unique = [&dim_set](int64 i) -> bool { in ValidateDotDimensionNumbers()631 return dim_set.insert(i).second; in ValidateDotDimensionNumbers()