#13646: Add second dimension support in binary op

tenstorrent · Oct 22, 2024 · 5e3cf35 · 5e3cf35
1 parent 204e3aa
commit 5e3cf35
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/ttnn/cpp/ttnn/operations/eltwise/binary/binary.cpp b/ttnn/cpp/ttnn/operations/eltwise/binary/binary.cpp
@@ -117,6 +117,12 @@ auto preprocess_inputs(
             Shape repeats(std::array<uint32_t, 4>{first_shape[0], 1, 1, 1});
             second = ttnn::repeat(second, repeats);
         }
+        // repeats second if it is smaller
+        if (first_shape.rank() == 4 and second_shape.rank() == 4 and first_shape[1] > second_shape[1]) {
+            tt::log_warning(tt::LogOp, "Using repeat op to broadcast channel dim");
+            Shape repeats(std::array<uint32_t, 4>{1, first_shape[1], 1, 1});
+            second = ttnn::repeat(second, repeats);
+        }
     };
     repeat_smaller(input_tensor_a, input_tensor_b);
     repeat_smaller(input_tensor_b, input_tensor_a);