Repository: tenstorrent/tt-metal
Pull request #20500: EXTERNAL PR: TTNN Support for stack [adding new ttnn op]
Author: login: jennychristopher
Created: 2025-04-10T10:11:34Z | Merged: 2025-04-24T01:47:14Z

### N/A
Link to Github Issue : N/A

### Problem description
ttnn doesn't have support for stack

### What's changed
Enabled support for stack op in ttnn as a composite op and added pytest

### Checklist
- [ ] [All post commit](https://github.com/tenstorrent/tt-metal/actions/workflows/all-post-commit-workflows.yaml) CI passes
- [ ] [Blackhole Post commit](https://github.com/tenstorrent/tt-metal/actions/workflows/blackhole-post-commit.yaml) CI passes (if applicable)
- [ ] [Model regression](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-models.yaml) CI passes (if applicable)
- [ ] [Device performance regression](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-device-models.yaml) CI passes (if applicable)
- [ ] **(For models and ops writers)** Full [new models tests](https://github.com/tenstorrent/tt-metal/actions/workflows/full-new-models-suite.yaml) CI passes (if applicable)
- [ ] New/Existing tests provide coverage for changes


Diff:
diff --git a/tests/ttnn/unit_tests/operations/test_stack.py b/tests/ttnn/unit_tests/operations/test_stack.py
new file mode 100755
index 000000000000..fc24d7aa6ed9
--- /dev/null
+++ b/tests/ttnn/unit_tests/operations/test_stack.py
@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+import ttnn
+from tests.ttnn.utils_for_testing import assert_with_pcc
+
+
+@pytest.mark.parametrize(
+    "input_shapes, dim, layout",
+    [
+        ([(1, 1, 253), (1, 1, 253), (1, 1, 253)], 2, ttnn.ROW_MAJOR_LAYOUT),
+        ([(1, 253), (1, 253), (1, 253)], 1, ttnn.ROW_MAJOR_LAYOUT),
+        ([(57, 83), (57, 83), (57, 83)], 0, ttnn.TILE_LAYOUT),
+        ([(8732,), (8732,), (8732,)], 0, ttnn.ROW_MAJOR_LAYOUT),
+        ([(123, 259), (123, 259), (123, 259)], -1, ttnn.TILE_LAYOUT),
+        ([(1, 1, 253), (1, 1, 253), (1, 1, 253)], -3, ttnn.ROW_MAJOR_LAYOUT),
+        ([(120, 1, 253, 3), (120, 1, 253, 3), (120, 1, 253, 3), (120, 1, 253, 3)], 3, ttnn.ROW_MAJOR_LAYOUT),
+        ([(57, 83), (57, 83), (57, 83), (57, 83), (57, 83), (57, 83), (57, 83)], -1, ttnn.TILE_LAYOUT),
+        ([(24, 50, 83), (24, 50, 83), (24, 50, 83), (24, 50, 83), (24, 50, 83), (24, 50, 83)], 2, ttnn.TILE_LAYOUT),
+    ],
+)
+def test_stack(device, input_shapes, dim, layout):
+    torch_tensors = [torch.rand(shape, dtype=torch.bfloat16) for shape in input_shapes]
+
+    ttnn_tensors = [ttnn.from_torch(tensor) for tensor in torch_tensors]
+    ttnn_tensors = [ttnn.to_device(tensor, device, memory_config=ttnn.L1_MEMORY_CONFIG) for tensor in ttnn_tensors]
+    ttnn_tensors = [ttnn.to_layout(tensor, layout) for tensor in ttnn_tensors]
+    torch_result = torch.stack(torch_tensors, dim=dim)
+    ttnn_result = ttnn.stack(ttnn_tensors, dim=dim)
+
+    ttnn_result_torch = ttnn.to_torch(ttnn_result)
+    assert_with_pcc(torch_result, ttnn_result_torch)
+
+    assert torch_result.shape == ttnn_result_torch.shape
diff --git a/ttnn/CMakeLists.txt b/ttnn/CMakeLists.txt
old mode 100644
new mode 100755
index 2004f68bc011..9d72379d6780
--- a/ttnn/CMakeLists.txt
+++ b/ttnn/CMakeLists.txt
@@ -352,6 +352,7 @@ set(TTNN_OP_SRCS
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/reshape_view/device/host/reshape_tiled_program_factory.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/unsqueeze/unsqueeze.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/squeeze/squeeze.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/stack/stack.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/sharded_partial/interleaved_to_sharded_partial/interleaved_to_sharded_partial.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/sharded_partial/interleaved_to_sharded_partial/device/interleaved_to_sharded_partial_op.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/sharded_partial/sharded_to_interleaved_partial/sharded_to_interleaved_partial.cpp
@@ -888,6 +889,7 @@ set(TTNN_SRC_PYBIND
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/repeat/repeat_pybind.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/reshape_view/reshape_pybind.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/reshape_on_device/reshape_pybind.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/stack/stack_pybind.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/sharded/interleaved_to_sharded/interleaved_to_sharded_pybind.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/sharded/reshard/reshard_pybind.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/data_movement/sharded/sharded_to_interleaved/sharded_to_interleaved_pybind.cpp
diff --git a/ttnn/cpp/ttnn/operations/data_movement/data_movement_pybind.hpp b/ttnn/cpp/ttnn/operations/data_movement/data_movement_pybind.hpp
old mode 100644
new mode 100755
index 8bce4f6f3877..7262cb2e7566
--- a/ttnn/cpp/ttnn/operations/data_movement/data_movement_pybind.hpp
+++ b/ttnn/cpp/ttnn/operations/data_movement/data_movement_pybind.hpp
@@ -34,6 +34,7 @@
 #include "ttnn/operations/data_movement/slice/slice_pybind.hpp"
 #include "ttnn/operations/data_movement/split/split_pybind.hpp"
 #include "ttnn/operations/data_movement/squeeze/squeeze_pybind.hpp"
+#include "ttnn/operations/data_movement/stack/stack_pybind.hpp"
 #include "ttnn/operations/data_movement/tilize/tilize_pybind.hpp"
 #include "ttnn/operations/data_movement/tilize_with_val_padding/tilize_with_val_padding_pybind.hpp"
 #include "ttnn/operations/data_movement/transpose/transpose_pybind.hpp"
@@ -81,6 +82,7 @@ void py_module(py::module& module) {
     py_bind_sharded_to_interleaved(module);
     py_bind_sharded_to_interleaved_partial(module);
     py_bind_squeeze(module);
+    py_bind_stack(module);
     py_bind_unsqueeze(module);
 }
 }  // namespace data_movement
diff --git a/ttnn/cpp/ttnn/operations/data_movement/stack/stack.cpp b/ttnn/cpp/ttnn/operations/data_movement/stack/stack.cpp
new file mode 100644
index 000000000000..c1070f886fb4
--- /dev/null
+++ b/ttnn/cpp/ttnn/operations/data_movement/stack/stack.cpp
@@ -0,0 +1,21 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "stack.hpp"
+#include "ttnn/operations/core/core.hpp"
+#include "cpp/ttnn/operations/data_movement/unsqueeze/unsqueeze.hpp"
+#include "cpp/ttnn/operations/data_movement/concat/concat.hpp"
+
+namespace ttnn::operations::data_movement {
+
+ttnn::Tensor StackOperation::invoke(const std::vector<ttnn::Tensor>& input_tensors, const int dim) {
+    TT_FATAL(!input_tensors.empty(), "Stack expects at least one tensor");
+    std::vector<ttnn::Tensor> expanded_tensors;
+    for (const auto& tensor : input_tensors) {
+        expanded_tensors.push_back(ttnn::unsqueeze(tensor, dim));
+    }
+    return ttnn::concat(expanded_tensors, dim);
+}
+
+}  // namespace ttnn::operations::data_movement
diff --git a/ttnn/cpp/ttnn/operations/data_movement/stack/stack.hpp b/ttnn/cpp/ttnn/operations/data_movement/stack/stack.hpp
new file mode 100644
index 000000000000..db08bcc24986
--- /dev/null
+++ b/ttnn/cpp/ttnn/operations/data_movement/stack/stack.hpp
@@ -0,0 +1,20 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "ttnn/decorators.hpp"
+
+namespace ttnn {
+namespace operations::data_movement {
+
+struct StackOperation {
+    static ttnn::Tensor invoke(const std::vector<ttnn::Tensor>& input_tensors, const int dim);
+};
+
+}  // namespace operations::data_movement
+
+constexpr auto stack = ttnn::register_operation<"ttnn::stack", ttnn::operations::data_movement::StackOperation>();
+
+}  // namespace ttnn
diff --git a/ttnn/cpp/ttnn/operations/data_movement/stack/stack_pybind.cpp b/ttnn/cpp/ttnn/operations/data_movement/stack/stack_pybind.cpp
new file mode 100644
index 000000000000..66cb92ca5899
--- /dev/null
+++ b/ttnn/cpp/ttnn/operations/data_movement/stack/stack_pybind.cpp
@@ -0,0 +1,52 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "stack_pybind.hpp"
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "pybind11/decorators.hpp"
+#include "ttnn/operations/data_movement/stack/stack.hpp"
+#include "ttnn/types.hpp"
+
+namespace ttnn::operations::data_movement {
+
+namespace detail {
+
+template <typename data_movement_operation_t>
+void bind_stack(pybind11::module& module, const data_movement_operation_t& operation, const char* doc) {
+    bind_registered_operation(
+        module,
+        operation,
+        doc,
+        ttnn::pybind_overload_t{
+            [](const data_movement_operation_t& self, const std::vector<ttnn::Tensor>& input_tensors, const int dim)
+                -> ttnn::Tensor { return self(input_tensors, dim); },
+            py::arg("input_tensors"),
+            py::arg("dim")});
+}
+
+}  // namespace detail
+
+void py_bind_stack(pybind11::module& module) {
+    detail::bind_stack(
+        module,
+        ttnn::stack,
+        R"doc(stack(input_tensors: List[ttnn.Tensor], dim: int) -> ttnn.Tensor
+
+        Stacks tensors along a new dimension.
+
+        Args:
+            * :attr:`input_tensors`: List of tensors to stack.
+            * :attr:`dim`: Dimension along which to stack.
+
+        Example:
+           >>> input_tensor = ttnn.from_torch(torch.randn((2, 2), dtype=torch.bfloat16), device=device)
+           >>> output = ttnn.stack((input_tensor,input_tensor), 1)
+
+        )doc");
+}
+
+}  // namespace ttnn::operations::data_movement
diff --git a/ttnn/cpp/ttnn/operations/data_movement/stack/stack_pybind.hpp b/ttnn/cpp/ttnn/operations/data_movement/stack/stack_pybind.hpp
new file mode 100644
index 000000000000..6c1befb8d2eb
--- /dev/null
+++ b/ttnn/cpp/ttnn/operations/data_movement/stack/stack_pybind.hpp
@@ -0,0 +1,13 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "pybind11/pybind_fwd.hpp"
+
+namespace ttnn::operations::data_movement {
+
+void py_bind_stack(pybind11::module& module);
+
+}  // namespace ttnn::operations::data_movement


Reviews:
user: login: ntarafdar
body: Great job
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T03:39:21Z
user: login: bbradelTT
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T13:44:31Z
user: login: ayerofieiev-tt
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-21T16:11:45Z
user: login: ayerofieiev-tt
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-21T16:15:06Z
user: login: jennychristopher
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T05:02:36Z

Review comments:
diff_hunk: @@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+import ttnn
+from tests.ttnn.utils_for_testing import assert_with_pcc
+
+
+@pytest.mark.parametrize(
+    "input_shapes, dim, layout",
+    [
+        ([(1, 1, 253), (1, 1, 253), (1, 1, 253)], 2, ttnn.ROW_MAJOR_LAYOUT),
path: tests/ttnn/unit_tests/operations/test_stack.py
user: login: ayerofieiev-tt
body: does it work to stack multiple tensors which require any kind of broadcasting?
created_at: 2025-04-21T16:11:45Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+import ttnn
+from tests.ttnn.utils_for_testing import assert_with_pcc
+
+
+@pytest.mark.parametrize(
+    "input_shapes, dim, layout",
+    [
+        ([(1, 1, 253), (1, 1, 253), (1, 1, 253)], 2, ttnn.ROW_MAJOR_LAYOUT),
path: tests/ttnn/unit_tests/operations/test_stack.py
user: login: jennychristopher
body: > does it work to stack multiple tensors which require any kind of broadcasting?

@ayerofieiev-tt It can stack tensors only of same shape so it does not require broadcasting.
created_at: 2025-04-22T05:02:36Z
author_association: CONTRIBUTOR

PR comments:
user: login: jennychristopher
created_at: 2025-04-10T10:14:54Z
body: @tt-rkim I have created this new PR for stack op and have addressed all your review comments that you have added in our old PR for stack - https://github.com/tenstorrent/tt-metal/pull/19411. Could you kindly review this ?
author_association: CONTRIBUTOR
user: login: jennychristopher
created_at: 2025-04-11T04:50:41Z
body: > Great job

Thanks @ntarafdar 
author_association: CONTRIBUTOR
user: login: jennychristopher
created_at: 2025-04-11T13:28:12Z
body: Can we go for CI runs ?
author_association: CONTRIBUTOR
user: login: ntarafdar
created_at: 2025-04-11T14:28:54Z
body: @jennychristopher yes can do! can you let me know after you rebase and I'll run them. 
author_association: CONTRIBUTOR
user: login: jennychristopher
created_at: 2025-04-11T15:28:54Z
body: > @jennychristopher yes can do! can you let me know after you rebase and I'll run them.

@ntarafdar rebased
author_association: CONTRIBUTOR
user: login: jennychristopher
created_at: 2025-04-14T05:23:13Z
body: @ntarafdar @tt-rkim I have rebased . Can we go for CI runs ?
author_association: CONTRIBUTOR
user: login: tt-rkim
created_at: 2025-04-14T12:18:52Z
body: @kkabilarTT Please assist with running APC, notice it's a ttnn unit test
author_association: COLLABORATOR
user: login: kkabilarTT
created_at: 2025-04-14T14:18:38Z
body: APC ran and passed here: https://github.com/tenstorrent/tt-metal/actions/runs/14447975097
author_association: CONTRIBUTOR
user: login: jennychristopher
created_at: 2025-04-16T05:02:37Z
body: @tt-rkim APC runs got passed, how do we proceed next. Do we have some other runs as well ?
author_association: CONTRIBUTOR
user: login: tt-rkim
created_at: 2025-04-20T18:42:58Z
body: <img width="910" alt="Screenshot 2025-04-20 at 2 42 29 PM" src="https://github.com/user-attachments/assets/3acb54d9-146a-46ef-a00e-17b12c168055" />

Running APC now.

We need to get your PR approved. Do you have access to our slack to get requests? Otherwise please get MCW teammates to help you out to ping necessary people to approve.
author_association: COLLABORATOR
user: login: tt-rkim
created_at: 2025-04-20T18:43:17Z
body: <img width="745" alt="Screenshot 2025-04-20 at 2 43 09 PM" src="https://github.com/user-attachments/assets/5217a460-47a9-4356-a637-f1f7a9ac46b4" />

Looks like these 3 need to help
author_association: COLLABORATOR
user: login: jennychristopher
created_at: 2025-04-21T04:10:35Z
body: > <img alt="Screenshot 2025-04-20 at 2 42 29 PM" width="910" src="https://private-user-images.githubusercontent.com/109366641/435490313-3acb54d9-146a-46ef-a00e-17b12c168055.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NDUyMDg4MzUsIm5iZiI6MTc0NTIwODUzNSwicGF0aCI6Ii8xMDkzNjY2NDEvNDM1NDkwMzEzLTNhY2I1NGQ5LTE0NmEtNDZlZi1hMDBlLTE3YjEyYzE2ODA1NS5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjUwNDIxJTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MDQyMVQwNDA4NTVaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT01NWMyZDUwNTg2MDJmOTRmNjhhN2E1NWNmOGU5MTQ3YThjOTFlMGE4ODRlZTk5NDA4NjEyMmFhNGQ3Y2QzNDNiJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.CIjGaAbCs3eSSplH2LF05xKK9BEZzWNqUs2DBK4qV5c">
> Running APC now.
> 
> We need to get your PR approved. Do you have access to our slack to get requests? Otherwise please get MCW teammates to help you out to ping necessary people to approve.

@tt-rkim I don't have access to your slack. Will check with my MCW teammates. Thanks !
author_association: CONTRIBUTOR
user: login: tt-rkim
created_at: 2025-04-21T11:38:45Z
body: Did you rebase go okay? Doesn't seem like it. Please retry and ensure only your changes are rebased onto latest main.
author_association: COLLABORATOR
user: login: tt-rkim
created_at: 2025-04-21T15:33:12Z
body: Workflows are running
author_association: COLLABORATOR
user: login: tt-rkim
created_at: 2025-04-22T17:34:02Z
body: Running CI. A rebase or merge with main will invalidate the run.
author_association: COLLABORATOR
user: login: tt-rkim
created_at: 2025-04-22T17:34:19Z
body: Please check when it's passing. Note you can check the statuses below, the Xs and checkmarks.
author_association: COLLABORATOR
user: login: bbradelTT
created_at: 2025-04-22T20:45:00Z
body: > Please check when it's passing. Note you can check the statuses below, the Xs and checkmarks.

@tt-rkim APC failed with what looks like some kind of infra error message https://github.com/tenstorrent/tt-metal/actions/runs/14600493431/job/40959652019?pr=20500
author_association: CONTRIBUTOR
user: login: tt-rkim
created_at: 2025-04-24T00:55:59Z
body: I'm not sure what happened there. Looks like something required building a new image, but I can't quite see from this [page ](https://github.com/tenstorrent/tt-metal/pkgs/container/tt-metal%2Ftt-metalium%2Fubuntu-24.04-ci-build-amd64/400353054?tag=5c818b6852c8b13a11eb670937ac5fca7aa51a43)when it was published.

The reason it tried to do this is because it detected that the image [didn't exist](https://github.com/tenstorrent/tt-metal/actions/runs/14600493431/job/40959648577?pr=20500). Nothing in this PR should've triggered it... so I'm curious exactly what happened.

Anyway, it shouldn't block as the other tests went through. What's important are the ttnn unit tests.

I'll approve it one last time so it can be merged.
author_association: COLLABORATOR
user: login: tt-rkim
created_at: 2025-04-24T00:56:26Z
body: Please click "Merge when ready" for yourself next time.
author_association: COLLABORATOR
user: login: jennychristopher
created_at: 2025-04-24T01:05:59Z
body: > Please click "Merge when ready" for yourself next time.

Sure. Will do
author_association: CONTRIBUTOR