Repository: tinygrad/tinygrad

Issue:
Windows Stack Overflow on EfficentNet Test `_deepwalk`
I get a windows stack overflow with
- `Python 3.9.0`
- `python .\test\models\test_train.py TestTrain.test_efficientnet`

With the following trace if I enable `faulthandler` (note the first error is just related to tqdm and can be ignored).
```
test\models\test_train.py .Windows fatal exception: stack overflow

Thread 0x00002074 (most recent call first):
  File "C:\Users\Patrick\AppData\Local\Programs\Python\Python39\lib\threading.py", line 316 in wait
  File "C:\Users\Patrick\AppData\Local\Programs\Python\Python39\lib\threading.py", line 574 in wait
  File "D:\COMMAI\tinygrad\venv\lib\site-packages\tqdm\_monitor.py", line 60 in run
  File "C:\Users\Patrick\AppData\Local\Programs\Python\Python39\lib\threading.py", line 950 in _bootstrap_inner
  File "C:\Users\Patrick\AppData\Local\Programs\Python\Python39\lib\threading.py", line 908 in _bootstrap

Current thread 0x00002238 (most recent call first):
  File "D:\COMMAI\tinygrad\tinygrad\tensor.py", line 159 in <listcomp>
  File "D:\COMMAI\tinygrad\tinygrad\tensor.py", line 159 in _deepwalk
  File "D:\COMMAI\tinygrad\tinygrad\tensor.py", line 159 in <listcomp>
  File "D:\COMMAI\tinygrad\tinygrad\tensor.py", line 159 in _deepwalk
  File "D:\COMMAI\tinygrad\tinygrad\tensor.py", line 159 in <listcomp>
  File "D:\COMMAI\tinygrad\tinygrad\tensor.py", line 159 in _deepwalk
  File "D:\COMMAI\tinygrad\tinygrad\tensor.py", line 159 in <listcomp>
...
```

If I change the deepwalk loop in [tensor.py](https://github.com/geohot/tinygrad/blob/b02a392d69c9180b789140b47aa861c57e239bcd/tinygrad/tensor.py#L155-L162) from this:
```python
if node._ctx:
  [_deepwalk(i, visited, nodes) for i in node._ctx.parents if i not in visited]
```
To this:
```python
for i in node._ctx.parents:
  if i not in visited:
    _deepwalk(i, visited, nodes)
```
Then all what I get is just the following warning and the test passes:
```
:\commai\tinygrad\tinygrad\ops.py:93: RuntimeWarning: overflow encountered in multiply
  else: ret = cls(cls.fxn_for_op[ast.op](*([x._buf for x in srcs] + ([ast.arg] if ast.arg else []))))
D:\COMMAI\tinygrad\venv\lib\site-packages\numpy\core\_methods.py:49: RuntimeWarning: overflow encountered in reduce
  return umr_sum(a, axis, dtype, out, keepdims, initial, where)
```

Not sure why that would make a difference (not that familiar with python inline loops). Any thoughts?


Discussion:
Hmm, maybe tail call recursion. Even though it's two lines, I'd merge the PR to change. It's currently building a list it doesn't need also.
This works:
```python
for i in [k for k in node._ctx.parents if k not in visited]:
  _deepwalk(i, visited, nodes)
```

but this does not:
```python
to_visit = [k for k in node._ctx.parents if k not in visited]
[_deepwalk(i, visited, nodes) for i in to_visit]
```


Resolving patch:
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -156,7 +156,9 @@ def deepwalk(self):
     def _deepwalk(node, visited, nodes):
       visited.add(node)
       if node._ctx:
-        [_deepwalk(i, visited, nodes) for i in node._ctx.parents if i not in visited]
+        for i in node._ctx.parents:
+          if i not in visited:
+            _deepwalk(i, visited, nodes)
         nodes.append(node)
       return nodes
     return _deepwalk(self, set(), [])


Test patch: