remove bitcast backward in function.py (tinygrad#7031)

bitcast cannot backward
chenyuxyz · Oct 13, 2024 · 13575f0 · 13575f0
1 parent ace834e
commit 13575f0
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/test/test_ops.py b/test/test_ops.py
@@ -2144,6 +2144,9 @@ def test_cast(self):
     helper_test_op([(3, 3)], lambda x: x.int(), forward_only=True)
     helper_test_op([(3, 3)], lambda x: x.bool(), forward_only=True)
 
+  def test_bitcast(self):
+    helper_test_op([(3, 3)], lambda x: x.view(torch.int32), lambda x: x.bitcast(dtypes.int32), forward_only=True)
+
 class TestOpsUint8(unittest.TestCase):
   @unittest.skip('this is broken for negative numbers')
   def test_cast(self):

diff --git a/tinygrad/function.py b/tinygrad/function.py
@@ -21,7 +21,8 @@ def forward(self, x:LazyBuffer, dtype:DType, bitcast:bool=False) -> LazyBuffer:
     return x.bitcast(dtype) if self.bitcast else x.cast(dtype)
 
   def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    return grad_output.bitcast(self.input_dtype) if self.bitcast else grad_output.cast(self.input_dtype)
+    if self.bitcast: raise RuntimeError("bitcast cannot backward")
+    return grad_output.cast(self.input_dtype)
 
 # ************* unary ops *************