Repository: tenstorrent/tt-forge-fe
Pull request #1416: Add conv2d_transpose support
Author: login: mstojkovicTT
Created: 2025-03-10T10:51:45Z | Merged: 2025-03-22T15:40:43Z

### Ticket
fixes: #1415 

### Problem description
Found Unsupported operations while lowering from TTForge to TTIR in forward graph

### What's changed
Added support for conv2d_tranpose op as well as e2e tests



Diff:
diff --git a/forge/csrc/passes/lower_to_mlir.cpp b/forge/csrc/passes/lower_to_mlir.cpp
index 501e5f6a74..bb6e0f55e1 100644
--- a/forge/csrc/passes/lower_to_mlir.cpp
+++ b/forge/csrc/passes/lower_to_mlir.cpp
@@ -4,10 +4,12 @@
 #include "lower_to_mlir.hpp"
 
 // Standard headers
+#include <cstddef>
 #include <cstdint>
 #include <stdexcept>
 #include <string>
 #include <utils/assert.hpp>
+#include <variant>
 
 // TTForge headers
 #include "forge_graph_module.hpp"
@@ -107,13 +109,31 @@ class AttributeMapper
     void initialize_default_mappings()
     {
         // Sort the mappings in lexicographical order
+
+        // conv2d_transpose
+        add_op_mapping("conv2d_transpose", "dilation", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
+        add_op_mapping("conv2d_transpose", "groups", AttributeRemap(std::nullopt, TargetType::I32Attr));
+        add_op_mapping(
+            "conv2d_transpose", "output_padding", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
+        add_op_mapping("conv2d_transpose", "padding", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
+        add_op_mapping("conv2d_transpose", "stride", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
+
+        // conv2d
         add_op_mapping("conv2d", "dilation", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
         add_op_mapping("conv2d", "groups", AttributeRemap(std::nullopt, TargetType::I32Attr));
         add_op_mapping("conv2d", "padding", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
         add_op_mapping("conv2d", "stride", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
+
+        // cumsum
         add_op_mapping("cumsum", "dim", AttributeRemap(std::nullopt, TargetType::I64Attr));
+
+        // reduce_avg
         add_op_mapping("reduce_avg", "dim", AttributeRemap("dim_arg"));
+
+        // repeat_interleave
         add_op_mapping("repeat_interleave", "repeats", AttributeRemap(std::nullopt, TargetType::UI32Attr));
+
+        // repeat
         add_op_mapping("repeat", "repeats", AttributeRemap("repeat_dimensions", TargetType::DenseI64ArrayAttr));
 
         // Add more default mappings here
@@ -654,6 +674,8 @@ class MLIRGenerator
         lowering_handler_map["repeat_interleave"] =
             &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::RepeatInterleaveOp>;
         lowering_handler_map["repeat"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::RepeatOp>;
+        lowering_handler_map["conv2d_transpose"] =
+            &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::ConvTranspose2dOp>;
         lowering_handler_map["reshape"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::ReshapeOp>;
         lowering_handler_map["select"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::SelectOp>;
         lowering_handler_map["sigmoid"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::SigmoidOp>;
diff --git a/forge/forge/op/convolution.py b/forge/forge/op/convolution.py
index 836ba9bda3..a77d83d1d4 100644
--- a/forge/forge/op/convolution.py
+++ b/forge/forge/op/convolution.py
@@ -76,6 +76,7 @@ def Conv2dTranspose(
     dilation: int = 1,
     groups: int = 1,
     channel_last: bool = False,
+    output_padding: Union[int, Tuple[int, int]] = 0,
 ) -> Tensor:
     """
     Conv2dTranspose transformation on input activations, with optional bias.
@@ -111,32 +112,22 @@ def Conv2dTranspose(
     if isinstance(padding, tuple):
         padding = list(padding)
 
+    if isinstance(dilation, int):
+        dilation = [dilation] * 2
+
     inputs = [activations, weights]
     if bias is not None:
         inputs.append(bias)
 
-    # Attrs are:
-    # [
-    #     stride_height,
-    #     stride_width,
-    #     dilation,
-    #     groups,
-    #     padding_left,
-    #     padding_right,
-    #     padding_top,
-    #     padding_bottom,
-    #     channel_last,
-    # ]
-    attrs = stride + [dilation, groups] + padding + [channel_last]
     return op(
         "conv2d_transpose",
         name,
         *inputs,
-        attrs=attrs,
         stride=stride,  # [sH, sW]
         dilation=dilation,  # [dH, dW]
         groups=groups,
         padding=padding,  # [pT, pL, pB, pR]
+        output_padding=output_padding,  # [opH, opW]
         channel_last=channel_last,
     ).get_tensor()
 
diff --git a/forge/forge/op/eval/forge/convolution.py b/forge/forge/op/eval/forge/convolution.py
index 006e11175f..165c4b3db9 100644
--- a/forge/forge/op/eval/forge/convolution.py
+++ b/forge/forge/op/eval/forge/convolution.py
@@ -189,7 +189,8 @@ def create(
         stride,  # Input format: [sH, sW]
         dilation,  # Input format: [dH, dW]
         groups,
-        padding,  # Input format: [pL, pR, pT, pB]
+        padding,  # Input format: [pT, pL, pB, pR]
+        output_padding,  # Input format: [opH, opW]
         channel_last,
     ):
         self = cls("conv2d_transpose")
@@ -197,6 +198,7 @@ def create(
         self.dilation = dilation
         self.groups = groups
         self.padding = padding
+        self.output_padding = output_padding
         self.channel_last = int(channel_last)
         return self
 
@@ -212,9 +214,9 @@ def eval(self, tensors):
         dilation = self.dilation
         groups = self.groups
         padding = (
-            self.padding[2],
             self.padding[0],
-        )  # [pT, pL] not sure why padding only has two elements (meenakshiramanathan1 PR #826)
+            self.padding[1],
+        )
 
         channel_last = self.channel_last
         if channel_last:
@@ -236,6 +238,7 @@ def eval(self, tensors):
             padding=padding,
             dilation=dilation,
             groups=groups,
+            output_padding=self.output_padding,
         )
 
         if channel_last:
@@ -248,10 +251,10 @@ def shape(self, tensor_shapes):
         stride_width = self.stride[1]
         dilation_height = self.dilation[0]
         dilation_width = self.dilation[1]
-        padding_left = self.padding[0]
-        padding_right = self.padding[1]
-        padding_top = self.padding[2]
-        padding_bottom = self.padding[3]
+        padding_top = self.padding[0]
+        padding_left = self.padding[1]
+        padding_bottom = self.padding[2]
+        padding_right = self.padding[3]
 
         act, weight = tensor_shapes[:2]
         batch_size = act[0]
@@ -260,8 +263,8 @@ def shape(self, tensor_shapes):
         h_in = act[-3] if self.channel_last else act[-2]
         w_in = act[-2] if self.channel_last else act[-1]
 
-        output_padding_height = 0
-        output_padding_width = 0
+        output_padding = self.output_padding
+        output_padding_height, output_padding_width = output_padding if output_padding else (0, 0)
 
         h_out = (
             (h_in - 1) * stride_height
@@ -319,11 +322,12 @@ def decompose(self, dc, inputs):
             new_inputs = [activations, weight] if bias is None else [activations, weight, bias]
             result = dc.op(
                 Conv2dTranspose.create(
-                    self.stride,
-                    self.dilation,
-                    self.groups,
-                    self.padding,
-                    True,  # If the original Conv2dTranspose was channel-last, that will not change.
+                    stride=self.stride,
+                    dilation=self.dilation,
+                    groups=self.groups,
+                    padding=self.padding,
+                    output_padding=self.output_padding,
+                    channel_last=True,  # If the original Conv2dTranspose was channel-last, that will not change.
                     # If it was channel-first, it the input will have been permuted by this point.
                     # So, the Conv2dTranspose op being created here is certainly channel-last.
                 ),
diff --git a/forge/forge/tvm_to_python.py b/forge/forge/tvm_to_python.py
index e2c558d5e0..44f1acab1b 100644
--- a/forge/forge/tvm_to_python.py
+++ b/forge/forge/tvm_to_python.py
@@ -706,6 +706,9 @@ def populate_conv2d_transpose_args(graph, nid, compiler_cfg):
     channel_last = int(node["attrs"]["data_layout"][0][0] == "NHWC")
     args.append(("channel_last", f"{channel_last}"))
 
+    output_padding = [int(opad) for opad in node["attrs"]["output_padding"][0]]
+    args.append(("output_padding", f"{output_padding}"))
+
     return args
 
 
diff --git a/forge/test/mlir/operators/nn/test_nn.py b/forge/test/mlir/operators/nn/test_nn.py
index 9ad3b47fb4..d5d99158f9 100644
--- a/forge/test/mlir/operators/nn/test_nn.py
+++ b/forge/test/mlir/operators/nn/test_nn.py
@@ -322,26 +322,225 @@ def forward(self, x):
     verify(inputs, framework_model, compiled_model, forge_property_handler=forge_property_recorder)
 
 
-@pytest.mark.xfail(reason="Found Unsupported operations while lowering from TTForge to TTIR in forward graph")
 @pytest.mark.parametrize(
     "in_channels, out_channels, kernel_size, stride, padding, groups, bias, dilation, padding_mode, input_shape",
     [
-        (16, 33, (3, 3), 2, 0, 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 32, (3, 5), 2, 1, 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 16, (3, 3), 1, 1, 16, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 1, (0, 0), 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 1, (1, 0), 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 1, (0, 1), 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 2, 0, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 33, (3, 3), 2, 0, 1, False, 1, "zeros", (20, 16, 50, 100)),
-        (16, 33, (3, 5), 2, 0, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 16, (5, 5), 1, 2, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 32, (3, 5), 2, 1, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 32, (3, 3), 4, (1, 2), 1, False, 1, "zeros", (20, 16, 50, 100)),
-        (16, 16, (3, 3), 2, (2, 3), 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 16, (3, 3), 1, (3, 3), 16, True, 1, "zeros", (20, 16, 50, 100)),
-        (64, 128, (7, 7), 4, (3, 5), 1, False, 1, "zeros", (16, 64, 80, 80)),
-        (32, 32, (1, 1), 1, (5, 6), 1, False, 1, "zeros", (10, 32, 20, 20)),
+        pytest.param(
+            16,
+            33,
+            (3, 3),
+            2,
+            0,
+            1,
+            True,
+            1,
+            "zeros",
+            (16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 2676352 B L1 buffer"),
+        ),
+        pytest.param(
+            16,
+            32,
+            (3, 5),
+            2,
+            1,
+            1,
+            True,
+            1,
+            "zeros",
+            (16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 2650240 B L1 buffer"),
+        ),
+        pytest.param(
+            16,
+            16,
+            (3, 3),
+            1,
+            1,
+            16,
+            True,
+            1,
+            "zeros",
+            (16, 50, 100),
+            marks=pytest.mark.xfail(
+                reason="Circular buffers clash with L1 buffers: static circular buffer region ends at 745248"
+            ),
+        ),
+        pytest.param(
+            16,
+            33,
+            (3, 3),
+            1,
+            (0, 0),
+            1,
+            True,
+            1,
+            "zeros",
+            (16, 50, 100),
+            marks=pytest.mark.xfail(
+                reason="Circular buffers clash with L1 buffers: static circular buffer region ends at 745248"
+            ),
+        ),
+        pytest.param(
+            16,
+            33,
+            (3, 3),
+            1,
+            (1, 0),
+            1,
+            True,
+            1,
+            "zeros",
+            (16, 50, 100),
+            marks=pytest.mark.xfail(
+                reason="Statically allocated circular buffers in program 33 clash with L1 buffers on core"
+            ),
+        ),
+        pytest.param(
+            16,
+            33,
+            (3, 3),
+            1,
+            (0, 1),
+            1,
+            True,
+            1,
+            "zeros",
+            (16, 50, 100),
+            marks=pytest.mark.xfail(
+                reason="Statically allocated circular buffers in program 39 clash with L1 buffers on core"
+            ),
+        ),
+        pytest.param(
+            16,
+            33,
+            (3, 3),
+            2,
+            0,
+            1,
+            True,
+            1,
+            "zeros",
+            (20, 16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 12800000 B L1 buffer"),
+        ),
+        pytest.param(
+            16,
+            33,
+            (3, 3),
+            2,
+            0,
+            1,
+            False,
+            1,
+            "zeros",
+            (20, 16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 12800000 B L1 buffer"),
+        ),
+        pytest.param(
+            16,
+            33,
+            (3, 5),
+            2,
+            0,
+            1,
+            True,
+            1,
+            "zeros",
+            (20, 16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 12800000 B L1 buffer"),
+        ),
+        pytest.param(
+            16,
+            16,
+            (5, 5),
+            1,
+            2,
+            1,
+            True,
+            1,
+            "zeros",
+            (20, 16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 12800000 B L1 buffer"),
+        ),
+        pytest.param(
+            16,
+            32,
+            (3, 5),
+            2,
+            1,
+            1,
+            True,
+            1,
+            "zeros",
+            (20, 16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 12800000 B L1 buffer"),
+        ),
+        pytest.param(
+            16,
+            32,
+            (3, 3),
+            4,
+            (1, 2),
+            1,
+            False,
+            1,
+            "zeros",
+            (20, 16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 12800000 B L1 buffer"),
+        ),
+        pytest.param(
+            16,
+            16,
+            (3, 3),
+            2,
+            (2, 3),
+            1,
+            True,
+            1,
+            "zeros",
+            (20, 16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 12800000 B L1 buffer"),
+        ),
+        pytest.param(
+            16,
+            16,
+            (3, 3),
+            1,
+            (3, 3),
+            16,
+            True,
+            1,
+            "zeros",
+            (20, 16, 50, 100),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 12800000 B L1 buffer"),
+        ),
+        pytest.param(
+            64,
+            128,
+            (7, 7),
+            4,
+            (3, 5),
+            1,
+            False,
+            1,
+            "zeros",
+            (16, 64, 80, 80),
+            marks=pytest.mark.xfail(reason="Out of Memory: Not enough space to allocate 26214400 B L1 buffer"),
+        ),
+        pytest.param(
+            32,
+            32,
+            (1, 1),
+            1,
+            (5, 6),
+            1,
+            False,
+            1,
+            "zeros",
+            (10, 32, 20, 20),
+            marks=pytest.mark.xfail(reason="Tensor mismatch. PCC = 0.0014"),
+        ),
     ],
 )
 def test_convtranspose2d(


Reviews:
user: login: nvukobratTT
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T04:54:58Z
user: login: mstojkovicTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T05:35:41Z
user: login: nvukobratTT
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-21T21:24:26Z

Review comments:
diff_hunk: @@ -107,6 +109,12 @@ class AttributeMapper
     void initialize_default_mappings()
     {
         // Sort the mappings in lexicographical order
+        add_op_mapping("conv2d_transpose", "dilation", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
path: forge/csrc/passes/lower_to_mlir.cpp
user: login: nvukobratTT
body: For readability, can we add block comments? E.g.
```cpp
        // Conv2d
        add_op_mapping("conv2d", "dilation", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
        add_op_mapping("conv2d", "groups", AttributeRemap(std::nullopt, TargetType::I32Attr));
        add_op_mapping("conv2d", "padding", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));

        // Conv2d Tranpose
        add_op_mapping("conv2d_transpose", "dilation", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
        add_op_mapping("conv2d_transpose", "groups", AttributeRemap(std::nullopt, TargetType::I32Attr));
        add_op_mapping(
            "conv2d_transpose", "output_padding", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
        add_op_mapping("conv2d_transpose", "padding", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
        add_op_mapping("conv2d_transpose", "stride", AttributeRemap(std::nullopt, TargetType::DenseI32ArrayAttr));
```
created_at: 2025-03-18T04:50:20Z
author_association: COLLABORATOR
diff_hunk: @@ -189,14 +189,16 @@ def create(
         stride,  # Input format: [sH, sW]
         dilation,  # Input format: [dH, dW]
         groups,
-        padding,  # Input format: [pL, pR, pT, pB]
+        padding,  # Input format: [pT, pL, pB, pR]
+        output_padding,
path: forge/forge/op/eval/forge/convolution.py
user: login: nvukobratTT
body: Can we add doc on the padding format? 
created_at: 2025-03-18T04:51:10Z
author_association: COLLABORATOR
diff_hunk: @@ -260,8 +263,13 @@ def shape(self, tensor_shapes):
         h_in = act[-3] if self.channel_last else act[-2]
         w_in = act[-2] if self.channel_last else act[-1]
 
-        output_padding_height = 0
-        output_padding_width = 0
+        output_padding = self.output_padding
+        if output_padding != 0:
path: forge/forge/op/eval/forge/convolution.py
user: login: nvukobratTT
body: To make this cleaner, can we do something like this:
```py
output_padding_height, output_padding_width = output_padding if output_padding else (0, 0)
```
created_at: 2025-03-18T04:53:05Z
author_association: COLLABORATOR
diff_hunk: @@ -301,29 +301,228 @@ def forward(self, x):
     verify(inputs, framework_model, compiled_model)
 
 
-@pytest.mark.xfail(reason="Found Unsupported operations while lowering from TTForge to TTIR in forward graph")
 @pytest.mark.parametrize(
     "in_channels, out_channels, kernel_size, stride, padding, groups, bias, dilation, padding_mode, input_shape",
     [
-        (16, 33, (3, 3), 2, 0, 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 32, (3, 5), 2, 1, 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 16, (3, 3), 1, 1, 16, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 1, (0, 0), 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 1, (1, 0), 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 1, (0, 1), 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 2, 0, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 33, (3, 3), 2, 0, 1, False, 1, "zeros", (20, 16, 50, 100)),
-        (16, 33, (3, 5), 2, 0, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 16, (5, 5), 1, 2, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 32, (3, 5), 2, 1, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 32, (3, 3), 4, (1, 2), 1, False, 1, "zeros", (20, 16, 50, 100)),
-        (16, 16, (3, 3), 2, (2, 3), 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 16, (3, 3), 1, (3, 3), 16, True, 1, "zeros", (20, 16, 50, 100)),
-        (64, 128, (7, 7), 4, (3, 5), 1, False, 1, "zeros", (16, 64, 80, 80)),
-        (32, 32, (1, 1), 1, (5, 6), 1, False, 1, "zeros", (10, 32, 20, 20)),
+        pytest.param(
path: forge/test/mlir/operators/nn/test_nn.py
user: login: nvukobratTT
body: Do we have any test case that actually passes? 
created_at: 2025-03-18T04:54:42Z
author_association: COLLABORATOR
diff_hunk: @@ -301,29 +301,228 @@ def forward(self, x):
     verify(inputs, framework_model, compiled_model)
 
 
-@pytest.mark.xfail(reason="Found Unsupported operations while lowering from TTForge to TTIR in forward graph")
 @pytest.mark.parametrize(
     "in_channels, out_channels, kernel_size, stride, padding, groups, bias, dilation, padding_mode, input_shape",
     [
-        (16, 33, (3, 3), 2, 0, 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 32, (3, 5), 2, 1, 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 16, (3, 3), 1, 1, 16, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 1, (0, 0), 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 1, (1, 0), 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 1, (0, 1), 1, True, 1, "zeros", (16, 50, 100)),
-        (16, 33, (3, 3), 2, 0, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 33, (3, 3), 2, 0, 1, False, 1, "zeros", (20, 16, 50, 100)),
-        (16, 33, (3, 5), 2, 0, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 16, (5, 5), 1, 2, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 32, (3, 5), 2, 1, 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 32, (3, 3), 4, (1, 2), 1, False, 1, "zeros", (20, 16, 50, 100)),
-        (16, 16, (3, 3), 2, (2, 3), 1, True, 1, "zeros", (20, 16, 50, 100)),
-        (16, 16, (3, 3), 1, (3, 3), 16, True, 1, "zeros", (20, 16, 50, 100)),
-        (64, 128, (7, 7), 4, (3, 5), 1, False, 1, "zeros", (16, 64, 80, 80)),
-        (32, 32, (1, 1), 1, (5, 6), 1, False, 1, "zeros", (10, 32, 20, 20)),
+        pytest.param(
path: forge/test/mlir/operators/nn/test_nn.py
user: login: mstojkovicTT
body: Not yet, it seems that all of them have the same problem with `Out of Memory` error
created_at: 2025-03-18T05:35:41Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov-commenter
created_at: 2025-03-10T11:41:37Z
body: ## [Codecov](https://app.codecov.io/gh/tenstorrent/tt-forge-fe/pull/1416?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 43.49%. Comparing base [(`5bb7546`)](https://app.codecov.io/gh/tenstorrent/tt-forge-fe/commit/5bb7546d2d3579bc6c716b5d8208ec762e575258?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent) to head [(`62b1c77`)](https://app.codecov.io/gh/tenstorrent/tt-forge-fe/commit/62b1c7733692accf1358b4232c211e59dc506b0c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main    #1416   +/-   ##
=======================================
  Coverage   43.49%   43.49%           
=======================================
  Files          49       49           
  Lines        7843     7843           
=======================================
  Hits         3411     3411           
  Misses       4432     4432           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/tenstorrent/tt-forge-fe/pull/1416?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).

author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-10T12:12:42Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>634 ran</td><td>492 passed</td><td>142 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-10T12:15:50Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>634 ran</td><td>492 passed</td><td>142 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-10T12:25:58Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>693 ran</td><td>556 passed</td><td>137 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-10T12:27:49Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>693 ran</td><td>556 passed</td><td>137 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T04:13:07Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>499 passed</td><td>141 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T04:20:31Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>499 passed</td><td>141 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T04:26:12Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T04:28:57Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T06:43:39Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>499 passed</td><td>141 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T06:47:20Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>499 passed</td><td>141 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T06:58:54Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T07:00:22Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T14:13:35Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>499 passed</td><td>141 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T14:16:33Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>499 passed</td><td>141 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T14:18:44Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-18T14:22:53Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T05:35:12Z
body: <table><tr><th><th>Tests</th><th>Passed ❌️</th><th>Skipped</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>0 ran</td><td>0 passed</td><td>0 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T05:50:09Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>641 ran</td><td>502 passed</td><td>139 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T05:51:38Z
body: <table><tr><th><th>Tests</th><th>Passed ☑️</th><th>Skipped ⚠️</th><th>Failed ❌️</th></tr><tr><td>TT-Forge-FE Tests</td><td>641 ran</td><td>501 passed</td><td>139 skipped</td><td>1 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2"><strong>TT-Forge-FE Tests</strong></td></tr><tr><td colspan="2"><em>pytest</em></td></tr><tr><td>test_llama_lm_head.test_llama_lm_head[openlm-research/open_llama_3b]</td><td>❌ failure</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T05:52:43Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>559 passed</td><td>140 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T07:37:14Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>499 passed</td><td>141 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T07:44:47Z
body: <table><tr><th><th>Tests</th><th>Passed ❌️</th><th>Skipped</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>0 ran</td><td>0 passed</td><td>0 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T07:49:40Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T07:51:33Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T15:31:48Z
body: <table><tr><th><th>Tests</th><th>Passed ❌️</th><th>Skipped</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>0 ran</td><td>0 passed</td><td>0 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-20T21:05:13Z
body: <table><tr><th><th>Tests</th><th>Passed ❌️</th><th>Skipped</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>0 ran</td><td>0 passed</td><td>0 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-21T18:11:47Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>492 passed</td><td>148 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-21T18:14:07Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>555 passed</td><td>144 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-21T18:19:01Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>492 passed</td><td>148 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-21T18:27:11Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>555 passed</td><td>144 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-21T22:20:03Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>492 passed</td><td>148 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-21T22:23:46Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>492 passed</td><td>148 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-21T22:25:58Z
body: <table><tr><th><th>Tests</th><th>Passed ❌️</th><th>Skipped</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>0 ran</td><td>0 passed</td><td>0 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-21T22:41:18Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>555 passed</td><td>144 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-21T23:24:56Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>555 passed</td><td>144 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE