#3188: refactor get_bcast_batch

tenstorrent · Nov 23, 2023 · ca97a68 · ca97a68
1 parent 16464c9
commit ca97a68
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/tests/tt_eager/python_api_testing/unit_testing/test_moreh_matmul.py b/tests/tt_eager/python_api_testing/unit_testing/test_moreh_matmul.py
@@ -67,6 +67,7 @@ def compare(tt_out, torch_out, atol=0.2, rtop=0.2):
     return allclose_result, isclose_true_ratio
 
 
+@skip_for_wormhole_b0
 @pytest.mark.parametrize(
     "input_a_shape",
     (
@@ -126,6 +127,7 @@ def test_moreh_matmul(input_a_shape, input_b_shape, device):
     assert allclose_result or isclose_true_ratio > 0.95
 
 
+@skip_for_wormhole_b0
 @pytest.mark.parametrize(
     "input_a_shape",
     (
@@ -190,6 +192,7 @@ def test_batched_moreh_matmul(input_a_shape, input_b_shape, device):
     assert allclose_result or isclose_true_ratio > 0.95
 
 
+@skip_for_wormhole_b0
 @pytest.mark.parametrize(
     "input_a_shape",
     (
@@ -250,6 +253,7 @@ def test_moreh_matmul_transpose_b(input_a_shape, input_b_shape, device):
     assert allclose_result or isclose_true_ratio > 0.95
 
 
+@skip_for_wormhole_b0
 @pytest.mark.parametrize(
     "input_a_shape",
     (
@@ -314,6 +318,7 @@ def test_batched_moreh_matmul_transpose_b(input_a_shape, input_b_shape, device):
     assert allclose_result or isclose_true_ratio > 0.95
 
 
+@skip_for_wormhole_b0
 @pytest.mark.parametrize(
     "input_shape",
     (

diff --git a/tt_eager/tt_dnn/op_library/moreh_matmul/multi_core/moreh_matmul_op_multi_core.cpp b/tt_eager/tt_dnn/op_library/moreh_matmul/multi_core/moreh_matmul_op_multi_core.cpp
@@ -17,9 +17,7 @@ namespace operations {
 namespace primary {
 
 std::tuple<bool, bool> get_bcast_batch(const Shape &input0_shape, const Shape &input1_shape) {
-    bool in0_bcast_batch = input0_shape[1] < input1_shape[1] ? true : false;
-    bool in1_bcast_batch = input0_shape[1] > input1_shape[1] ? true : false;
-    return {in0_bcast_batch, in1_bcast_batch};
+    return {(input0_shape[1] < input1_shape[1]), (input0_shape[1] > input1_shape[1])};
 }
 
 operation::ProgramWithCallbacks moreh_matmul_multi_core(