From a73b59caa2ed3521b2b00f19f46491d399b3c68e Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Tue, 14 Oct 2025 16:50:43 +0800
Subject: [PATCH] work on shape property

---
 tinygrad/uop/ops.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index cd18db40bf..6f0926829e 100644
--- a/tinygrad/uop/ops.py
+++ b/tinygrad/uop/ops.py
@@ -225,6 +225,17 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   @property
   def shape(self) -> tuple[sint, ...]:
+    match self.op:
+      case Ops.CONST: return ()
+      case Ops.BUFFER: return (self.arg,)
+      case Ops.BUFFER_VIEW: return (self.arg[0],)
+      case Ops.RESHAPE:
+        if prod(self.src[0].shape) != prod(self.arg): raise RuntimeError(f"bad reshape: {self.src[0].shape} -> {self.arg}")
+        return self.arg
+      case Ops.PERMUTE:
+        if sorted(self.arg) != list(range(len(bs:=self.src[0].shape))): raise RuntimeError(f"invalid permutation {self.arg} of len {len(bs)}")
+        return tuple(bs[i] for i in self.arg)
+      # TODO: finish this and remove self.st.shape
     assert self.st is not None, f"{self.op} doesn't have a shape"
     return unwrap(self.st).shape
   @property