Repository: tinygrad/tinygrad

Issue:
Bug: NaN when using `cat` on tensor with unrealised `log` op
Tries to take the log of zero due to merging of the ops, resulting in NAN values. 

To reproduce:

```
x0 = Tensor([1])
x1 = Tensor([1]).log()    # Can realize `x1` before calling `cat` to circumvent
print(x0.cat(x1).numpy()) # [nan  0.]
```

Kernel it generates:
```
#include <metal_stdlib>
using namespace metal;
kernel void E_2(device float* data0, const device float* data1, const device float* data2, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {
{ int gidx0 = gid.x;  /* 2 */
    float val1_0 = ((gidx0<1)) ? (float)(data1[0]) : 0.0f;
    float val2_0 = (((gidx0*-1)<0)) ? (float)(data2[0]) : 0.0f;
    float val3_0 = (((gidx0*-1)<0)) ? (float)(0.6931471805599453f) : 0.0f;
    float alu0 = log2(val2_0);
    float alu1 = (alu0*val3_0);
    float alu2 = (alu1+val1_0);
    data0[gidx0] = alu2;
  } /* global+local */
}
```


Discussion:

Resolving patch:
diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py
--- a/tinygrad/lazy.py
+++ b/tinygrad/lazy.py
@@ -283,8 +283,9 @@ def _push_movement_ops(srcs:Tuple[LazyBuffer, ...]) -> Tuple[LazyBuffer, ...]:
       assert isinstance(bx.op.op, MovementOps)
       mops.append((bx.op.op, bx.op.arg))
       bx = cast(LazyBuffer, bx.op.src[0])
-    # NOTE: can't push pads with a div
-    if not bx.realized and bx.optype == BinaryOps and len(bx.children) <= 1 and len(mops) and (all(x[0] != MovementOps.PAD for x in mops) or all(x.op != BinaryOps.DIV for x in bx.op.get_lazyops())):
+    # NOTE: can't push pads past anything where f(0, 0) != 0 or f(0) != 0
+    unsafe_pad_ops = {BinaryOps.DIV, BinaryOps.CMPEQ, UnaryOps.LOG2, UnaryOps.EXP2, UnaryOps.RECIP}
+    if not bx.realized and bx.optype == BinaryOps and len(bx.children) <= 1 and len(mops) and (all(x[0] != MovementOps.PAD for x in mops) or all(x.op not in unsafe_pad_ops for x in bx.op.get_lazyops())):
       new_srcs.append(bx.op.replace_with_movement_ops(mops[::-1]))
     else:
       new_srcs.append(x)


Test patch:
diff --git a/test/test_lazybuffer.py b/test/test_lazybuffer.py
--- a/test/test_lazybuffer.py
+++ b/test/test_lazybuffer.py
@@ -2,6 +2,7 @@
 import numpy as np
 import unittest
 from tinygrad.lazy import LazyBuffer
+from tinygrad.tensor import Tensor
 
 class TestLazyBuffer(unittest.TestCase):
   def test_fromcpu_buffer_sharing(self):
@@ -22,5 +23,25 @@ def helper(a: np.ndarray):
         for start in [0, 1]:
           helper(a[(slice(start, None, stride),)*ndims])
 
+  def test_shuffle_pad_ops_cmpeq(self):
+    y = Tensor([1]).cat(Tensor([1]).eq(0)).numpy()
+    z = Tensor([1, 0]).numpy()
+    np.testing.assert_allclose(y, z)
+
+  def test_shuffle_pad_ops_div(self):
+    y = Tensor([1]).cat(Tensor([1]).div(Tensor([2.0]))).numpy()
+    z = Tensor([1, 0.5]).numpy()
+    np.testing.assert_allclose(y, z)
+
+  def test_shuffle_pad_ops_log(self):
+    y = Tensor([1]).cat(Tensor([1]).log()).numpy()
+    z = Tensor([1, 0]).numpy()
+    np.testing.assert_allclose(y, z)
+
+  def test_shuffle_pad_ops_exp(self):
+    y = Tensor([1]).cat(Tensor([1]).exp()).numpy()
+    z = Tensor([1, np.e]).numpy()
+    np.testing.assert_allclose(y, z)
+
 if __name__ == "__main__":
   unittest.main()