From b67f75978049e02a04c2b9ce238e01cc4543bb06 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Fri, 12 Apr 2024 16:46:01 -0700
Subject: [PATCH] abstractions3 is currently wishful thinking (#4124)

* abstractions3 is currently wishful thinking

* a3

* work

* minor

* progress on a3

* more

* update abstractions3

* cleaner
---
 docs/abstractions3.py | 63 +++++++++++++++++++++++++++++++++++++++++++
 tinygrad/lazy.py      |  2 +-
 2 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 docs/abstractions3.py

diff --git a/docs/abstractions3.py b/docs/abstractions3.py
new file mode 100644
index 0000000000..a013afc7df
--- /dev/null
+++ b/docs/abstractions3.py
@@ -0,0 +1,63 @@
+# abstractions2 goes from back to front, here we will go from front to back
+from typing import List
+from tqdm import tqdm
+from tinygrad.helpers import DEBUG
+
+# *****
+# 0. Load mnist on the device
+
+from tinygrad.features.datasets import mnist
+X_train, Y_train, _, _ = mnist()
+X_train = X_train.float()
+X_train -= X_train.mean()
+
+# *****
+# 1. Define an MNIST model.
+
+from tinygrad import Tensor
+
+l1 = Tensor.kaiming_uniform(128, 784)
+l2 = Tensor.kaiming_uniform(10, 128)
+def model(x): return x.flatten(1).dot(l1.T).relu().dot(l2.T)
+l1n, l2n = l1.numpy(), l2.numpy()
+
+# *****
+# 2. Choose a batch for training and do the backward pass.
+
+from tinygrad.nn.optim import SGD
+optim = SGD([l1, l2])
+
+X, Y = X_train[samples:=Tensor.randint(128, high=X_train.shape[0])], Y_train[samples]
+optim.zero_grad()
+model(X).sparse_categorical_crossentropy(Y).backward()
+optim._step()   # this will step the optimizer without running realize
+
+# *****
+# 3. Create a schedule.
+
+# The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
+# l1.lazydata and l2.lazydata define a computation graph
+
+from tinygrad.ops import ScheduleItem
+from tinygrad.engine.schedule import create_schedule
+schedule: List[ScheduleItem] = create_schedule([l1.lazydata, l2.lazydata])
+
+print(f"The schedule contains {len(schedule)} items.")
+for si in schedule: print(str(si)[:80])
+
+# *****
+# 4. Lower a schedule.
+
+from tinygrad.engine.realize import lower_schedule_item, ExecItem
+lowered: List[ExecItem] = [ExecItem(lower_schedule_item(si), list(si.outputs+si.inputs)) for si in tqdm(schedule)]
+
+# *****
+# 5. Run the schedule
+
+for ei in tqdm(lowered): ei.run()
+
+# *****
+# 6. Print the weight change
+
+print("first weight change\n", l1.numpy()-l1n)
+print("second weight change\n", l2.numpy()-l2n)
diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py
index 8aaa23fdb0..b31a8dca1f 100644
--- a/tinygrad/lazy.py
+++ b/tinygrad/lazy.py
@@ -41,7 +41,7 @@ class LazyBuffer:
       self._base = base
 
   def __repr__(self) -> str:
-    return f"<LB {self.device} {self.shape} contig:{self.st.contiguous} {self.st if self.base != self else (self.op, self.realized)}>"
+    return f"<LB {self.device} {self.shape} {str(self.dtype)[7:]} {self.st if self.base != self else (self.op, self.realized)}>"
 
   @property
   def realized(self) -> Optional[Buffer]: