From c261c8cba76b1aa6c6a085b1e2ac5fb6f61ff6f9 Mon Sep 17 00:00:00 2001
From: Brian Liu <bliu@tenstorrent.com>
Date: Thu, 16 Jan 2025 22:25:51 +0000
Subject: [PATCH] #0: Fix for rank 0 and 1 tensors

---
 ttnn/cpp/ttnn/tensor/layout/tensor_layout.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/ttnn/cpp/ttnn/tensor/layout/tensor_layout.cpp b/ttnn/cpp/ttnn/tensor/layout/tensor_layout.cpp
index 5ba3da5c29d..057b4835ea9 100644
--- a/ttnn/cpp/ttnn/tensor/layout/tensor_layout.cpp
+++ b/ttnn/cpp/ttnn/tensor/layout/tensor_layout.cpp
@@ -262,7 +262,10 @@ Size TensorLayout::get_physical_shard_shape() const {
 
 Size TensorLayout::compute_logical_2d_shape(const ttnn::SimpleShape& shape) const {
     size_t width = shape[-1];
-    size_t height = shape.volume() / width;
+    size_t height = 1;
+    for (int i = -2; i >= -shape.rank(); --i) {
+        height *= shape[i];
+    }
     return Size{height, width};
 }