Add BroadcastIndexesRange tests with dims of size 1 in output

Per request from kimishpatel.

ghstack-source-id: 0624885cc1dac7ee43d1982c2011ded348ca4735
ghstack-comment-id: 2701662107
Pull Request resolved: #8964

Author: swolchok

kernels/portable/cpu/util/test/broadcast_indexes_range_test.cpp

@@ -112,8 +112,66 @@ TEST(BroadcastIndexesRangeTest, OneAndTwoDExhaustive) {
   EXPECT_EQ(expected, actual);
 }
 
-// Here we assume that the previous tests established that padding
-// with leading 1s is working, and test:
+// Make sure nothing is thrown off by a size-1 dim in the output:
+// [] -> [1, W]
+// [] -> [H, 1]
+// [1] -> [1, W]
+// [1] -> [H, 1]
+// [W] -> [1, W]
+// [1, 1] -> [1, W]
+// [1, 1] -> [H, 1]
+// [1, W] -> [1, W]
+// [H, 1] -> [H, 1]
+TEST(BroadcastIndexesRangeTest, OneAndTwoDWith1InOutputShapeExhaustive) {
+  TensorFactory<ScalarType::Int> tf;
+  constexpr auto H = 2;
+  constexpr auto W = 3;
+  Tensor out_row = tf.zeros({1, W});
+  Tensor out_col = tf.zeros({H, 1});
+  Tensor in_0d_scalar = tf.zeros({});
+  Tensor in_1d_scalar = tf.zeros({1});
+  Tensor in_2d_scalar = tf.zeros({1, 1});
+
+  Tensor in_row = tf.zeros({W});
+  Tensor in_leading_one_row = tf.zeros({1, W});
+
+  Tensor in_col = tf.zeros({H, 1});
+
+  size_t idx = 0;
+  for (const auto
+       [out_idx,
+        in_0d_idx,
+        in_1d_idx,
+        in_2d_idx,
+        in_row_idx,
+        in_leading_one_row_idx] :
+       BroadcastIndexesRange<5>(
+           out_row,
+           in_0d_scalar,
+           in_1d_scalar,
+           in_2d_scalar,
+           in_row,
+           in_leading_one_row)) {
+    EXPECT_EQ(out_idx, idx++);
+    EXPECT_EQ(in_0d_idx, 0);
+    EXPECT_EQ(in_1d_idx, 0);
+    EXPECT_EQ(in_2d_idx, 0);
+    EXPECT_EQ(in_row_idx, out_idx);
+    EXPECT_EQ(in_leading_one_row_idx, out_idx);
+  }
+
+  idx = 0;
+  for (const auto [out_idx, in_0d_idx, in_1d_idx, in_2d_idx, in_col_idx] :
+       BroadcastIndexesRange<4>(
+           out_col, in_0d_scalar, in_1d_scalar, in_2d_scalar, in_col)) {
+    EXPECT_EQ(out_idx, idx++);
+    EXPECT_EQ(in_0d_idx, 0);
+    EXPECT_EQ(in_1d_idx, 0);
+    EXPECT_EQ(in_2d_idx, 0);
+    EXPECT_EQ(in_col_idx, out_idx);
+  }
+}
+
 // [1, 1, 1] -> [C, H, W]
 // [C, H, 1] -> [C, H, W]
 // [C, 1, W] -> [C, H, W]
@@ -166,11 +224,12 @@ TEST(BroadcastIndexesRangeTest, ThreeDBroadcasting) {
 // 4-D should generalize, but we will go ahead and test:
 // [N, 1, H, 1] -> [N, C, H, W]
 // [1, C, 1, W] -> [N, C, H, W]
-TEST(BroadcastIndexesRangeTest, FourDBroadcasting) {
+template <size_t N, size_t C, size_t H, size_t W>
+void four_d_broadcasting_test() {
   TensorFactory<ScalarType::Int> tf;
-  Tensor out = tf.zeros({2, 3, 4, 5});
-  Tensor in_broadcast_cw = tf.zeros({2, 1, 4, 1});
-  Tensor in_broadcast_nh = tf.zeros({1, 3, 1, 5});
+  Tensor out = tf.zeros({N, C, H, W});
+  Tensor in_broadcast_cw = tf.zeros({N, 1, H, 1});
+  Tensor in_broadcast_nh = tf.zeros({1, C, 1, W});
 
   // Writing out all the indexes would be too cumbersome, so here we
   // take the opportunity to mutation test against delinearize_index
@@ -190,3 +249,12 @@ TEST(BroadcastIndexesRangeTest, FourDBroadcasting) {
         linearize_access_indexes(out_indexes, out.dim(), in_broadcast_nh));
   }
 }
+
+TEST(BroadcastIndexesRangeTest, FourDBroadcasting) {
+  four_d_broadcasting_test<2, 3, 4, 5>();
+}
+
+TEST(BroadcastIndexesRangeTest, FourDBroadcastingWithOneDimsInOutput) {
+  four_d_broadcasting_test<2, 3, 1, 5>();
+  four_d_broadcasting_test<2, 1, 3, 1>();
+}