document graph mode

README.md

@@ -119,7 +119,7 @@ Relu, Log, Exp                  # unary ops
 Sum, Max                        # reduce ops (with axis argument)
 Add, Sub, Mul, Pow              # binary ops (with broadcasting)
 Reshape, Transpose, Slice       # movement ops
-Matmul, Conv2D                  # processing ops
+Matmul, Conv2D(NCHW)            # processing ops
 ```
 
 While more ops may be added, I think this base is stable.
@@ -162,6 +162,13 @@ See `examples/yolov3.py`
 
 tinygrad will always be below 1000 lines. If it isn't, we will revert commits until tinygrad becomes smaller.
 
+### Drawing Execution Graph
+```bash
+GRAPH=1 python3 test/test_mnist.py TestMNIST.test_conv_onestep
+dot -Tsvg /tmp/net.dot -o /tmp/net.svg
+open /tmp/net.svg
+```
+
 ### Running tests
 
 ```bash

test/test_mnist.py

@@ -46,6 +46,11 @@ class TinyConvNet:
     return x.dot(self.l1).logsoftmax()
 
 class TestMNIST(unittest.TestCase):
+  def test_conv_onestep(self):
+    np.random.seed(1337)
+    model = TinyConvNet()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    train(model, X_train, Y_train, optimizer, steps=1)
 
   def test_conv(self):
     np.random.seed(1337)