window partition, reverse OCAB fix, ssr pcc: 0.999952

Author: ign-krishnanand

models/experimental/SSR/tests/test_ssr.py

@@ -155,7 +155,7 @@ def test_ssr_model(input_shape, num_cls, with_conv):
             )
 
         # Convert input to TTNN tensor
-        tt_input = ttnn.from_torch(x, device=device, layout=ttnn.TILE_LAYOUT)
+        tt_input = ttnn.from_torch(x, device=device, layout=ttnn.TILE_LAYOUT, dtype=ttnn.bfloat16)
 
         # Run TTNN model
         tt_sr, tt_patch_fea3 = tt_model(tt_input)

models/experimental/SSR/tt/OCAB.py

@@ -4,15 +4,35 @@
 
 
 def window_partition_ttnn(x, window_size):
-    """Partition into non-overlapping windows"""
-    B, H, W, C = x.shape
-    num_windows = (H // window_size) * (W // window_size)
-    return ttnn.reshape(x, [B * num_windows, window_size, window_size, C], memory_config=ttnn.L1_MEMORY_CONFIG)
+    """TTNN implementation of window partitioning"""
+    b, h, w, c = x.shape
 
+    # Reshape: (b, h, w, c) -> (b, h//ws, ws, w//ws, ws, c)
+    reshaped = ttnn.reshape(x, (b, h // window_size, window_size, w // window_size, window_size, c))
 
-def window_reverse_ttnn(windows, window_size, H, W):
-    B = windows.shape[0] // (H * W // window_size // window_size)
-    return ttnn.reshape(windows, [B, H, W, -1], memory_config=ttnn.L1_MEMORY_CONFIG)
+    # Permute: (0, 1, 3, 2, 4, 5) -> group windows together
+    permuted = ttnn.permute(reshaped, (0, 1, 3, 2, 4, 5))
+
+    # Final reshape to get windows
+    windows = ttnn.reshape(permuted, (-1, window_size, window_size, c))
+
+    return windows
+
+
+def window_reverse_ttnn(windows, window_size, h, w):
+    """TTNN implementation of window reverse"""
+    b = int(windows.shape[0] / (h * w / window_size / window_size))
+
+    # Reshape windows back to grid
+    reshaped = ttnn.reshape(windows, (b, h // window_size, w // window_size, window_size, window_size, -1))
+
+    # Permute back to original order
+    permuted = ttnn.permute(reshaped, (0, 1, 3, 2, 4, 5))
+
+    # Final reshape to original spatial dimensions
+    output = ttnn.reshape(permuted, (b, h, w, -1))
+
+    return output
 
 
 class TTOCAB(LightweightModule):