Repository: tinygrad/tinygrad

Issue:
Make REQUIRES_GRAD not the default to match PyTorch
It's annoying to write it everywhere, for most things you don't want the gradient. Will improve performance and require less no grads.


Discussion:

Resolving patch:
diff --git a/extra/utils.py b/extra/utils.py
--- a/extra/utils.py
+++ b/extra/utils.py
@@ -23,10 +23,11 @@ def fetch(url):
     os.rename(fp+".tmp", fp)
   return dat
 
+# TODO: move this to optim.py?
 def get_parameters(obj):
   parameters = []
   if isinstance(obj, Tensor):
-    if obj.requires_grad: parameters.append(obj)
+    parameters.append(obj)
   elif isinstance(obj, list) or isinstance(obj, tuple):
     for x in obj:
       parameters.extend(get_parameters(x))
diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py
--- a/tinygrad/nn/optim.py
+++ b/tinygrad/nn/optim.py
@@ -3,6 +3,11 @@
 
 class Optimizer:
   def __init__(self, params):
+    # if it's None, but being put into an optimizer, set it to True
+    for x in params:
+      if x.requires_grad is None:
+        x.requires_grad = True
+
     self.params = [x for x in params if x.requires_grad]
 
   def zero_grad(self):
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -14,7 +14,7 @@ class Tensor:
   # TODO: remove no_init when uniform is late bind
   training, no_grad, no_init = False, False, False
 
-  def __init__(self, data, device=Device.DEFAULT, requires_grad=True):
+  def __init__(self, data, device=Device.DEFAULT, requires_grad=None):
     if isinstance(data, list):
       data = np.array(data, dtype=np.float32)
     elif isinstance(data, LazyBuffer) and data.device != device:
@@ -32,7 +32,10 @@ def __init__(self, data, device=Device.DEFAULT, requires_grad=True):
 
     # tensors have gradients, buffers do not
     self.grad : Optional[Tensor] = None
-    self.requires_grad = requires_grad
+
+    # NOTE: this can be in three states. False and None: no gradient, True: gradient
+    # None (the default) will be updated to True if it's put in an optimizer
+    self.requires_grad : Optional[bool] = requires_grad
 
     # internal variables used for autograd graph construction
     self._ctx : Optional[Function] = None


Test patch:
diff --git a/test/test_gc.py b/test/test_gc.py
--- a/test/test_gc.py
+++ b/test/test_gc.py
@@ -9,22 +9,22 @@ def tensors_allocated():
 class TestGC(unittest.TestCase):
 
   def test_gc(self):
-    a = Tensor.zeros(4,4)
-    b = Tensor.zeros(4,4)
+    a = Tensor.zeros(4, 4, requires_grad=True)
+    b = Tensor.zeros(4, 4, requires_grad=True)
     (a*b).mean().backward()
     assert(tensors_allocated() > 0)
     del a,b
     assert(tensors_allocated() == 0)
 
   def test_gc_complex(self):
-    a = Tensor.zeros(4,4)
-    b = Tensor.zeros(4,4)
+    a = Tensor.zeros(4, 4, requires_grad=True)
+    b = Tensor.zeros(4, 4, requires_grad=True)
     assert(tensors_allocated() == 2)
     (a*b).mean().backward()
     assert(tensors_allocated() == 4)
     del b
     assert(tensors_allocated() == 2)
-    b = Tensor.zeros(4,4)
+    b = Tensor.zeros(4, 4, requires_grad=True)
     print(tensors_allocated())
     (a*b).mean().backward()
     print(tensors_allocated())
diff --git a/test/test_ops.py b/test/test_ops.py
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -14,7 +14,7 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn, atol=1e-6, rtol=1e-3, grad_ato
   else:
     ts = [torch.tensor((np.random.random(size=x).astype(np.float32)+a)*b, requires_grad=True) for x in shps]
 
-  tst = [Tensor(x.detach().numpy()) for x in ts]
+  tst = [Tensor(x.detach().numpy(), requires_grad=True) for x in ts]
   out = torch_fxn(*ts)
   ret = tinygrad_fxn(*tst)
 
diff --git a/test/test_tensor.py b/test/test_tensor.py
--- a/test/test_tensor.py
+++ b/test/test_tensor.py
@@ -22,9 +22,9 @@ def test_plus_equals(self):
 
   def test_backward_pass(self):
     def test_tinygrad():
-      x = Tensor(x_init)
-      W = Tensor(W_init)
-      m = Tensor(m_init, requires_grad=False)
+      x = Tensor(x_init, requires_grad=True)
+      W = Tensor(W_init, requires_grad=True)
+      m = Tensor(m_init)
       out = x.dot(W).relu()
       out = out.logsoftmax()
       out = out.mul(m).add(m).sum()
@@ -46,9 +46,9 @@ def test_pytorch():
 
   def test_backward_pass_diamond_model(self):
     def test_tinygrad():
-      u = Tensor(U_init)
-      v = Tensor(V_init)
-      w = Tensor(W_init)
+      u = Tensor(U_init, requires_grad=True)
+      v = Tensor(V_init, requires_grad=True)
+      w = Tensor(W_init, requires_grad=True)
       x = u.mul(v).relu()
       y = u.mul(w).relu()
       out = x.add(y).mul(y).relu()