simpler, but not inverse

test/null/test_bitcast_spec.py

@@ -0,0 +1,45 @@
+import unittest
+from tinygrad import UOp, dtypes
+from tinygrad.uop.ops import shape_to_shape_arg, ParamArg, Ops, AddrSpace
+
+def placeholder(shape, dtype, slot):
+  return UOp(Ops.PARAM, dtype, (shape_to_shape_arg(shape),), arg=ParamArg(slot, AddrSpace.GLOBAL))
+
+class TestBitcastSpec(unittest.TestCase):
+  def test_bitcast_no_shape_change(self):
+    pl = placeholder((10,10), dtypes.int, 0)
+    out = pl.bitcast(dtypes.float)
+    self.assertEqual(out.shape, (10,10))
+
+  def test_bitcast_increase_shape(self):
+    pl = placeholder((10,10), dtypes.int, 0)
+    out = pl.bitcast(dtypes.short)
+    self.assertEqual(out.shape, (10,20))
+
+  def test_bitcast_decrease_shape(self):
+    pl = placeholder((10,10), dtypes.int, 0)
+    out = pl.bitcast(dtypes.long)
+    self.assertEqual(out.shape, (10,5))
+
+  def test_bitcast_remove_ones(self):
+    pl = placeholder((10,2), dtypes.int, 0)
+    out = pl.bitcast(dtypes.long)
+    self.assertEqual(out.shape, (10,1))
+
+  def test_bitcast_remove_ones_full(self):
+    pl = placeholder((2,), dtypes.int, 0)
+    out = pl.bitcast(dtypes.long)
+    self.assertEqual(out.shape, (1,))
+
+  def test_bitcast_add_ones_full(self):
+    pl = placeholder((), dtypes.long, 0)
+    out = pl.bitcast(dtypes.int)
+    self.assertEqual(out.shape, (2,))
+
+  def test_bitcast_add_ones_full_uchar(self):
+    pl = placeholder((), dtypes.long, 0)
+    out = pl.bitcast(dtypes.uchar)
+    self.assertEqual(out.shape, (8,))
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad/uop/ops.py

@@ -312,10 +312,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
       # TODO: disallow shape changing bitcast
       case Ops.BITCAST:
-        ps = self.src[0]._shape
-        if ps is None: return None
+        ps = self.src[0].shape
         if (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize):
-          return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),) if len(ps) > 0 else ps
+          if ps == ():
+            if output_sz > input_sz: raise RuntimeError(f"shape {ps} must be an expanding bitcast {output_sz=} {input_sz=}")
+            ps = (1,)
+          return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),)
         return ps
 
       # MULTI marker has no shape