diff --git a/test/test_dtype.py b/test/test_dtype.py index 9440dab49b..fad8ed156b 100644 --- a/test/test_dtype.py +++ b/test/test_dtype.py @@ -137,16 +137,14 @@ class TestBFloat16(unittest.TestCase): assert tnp.dtype == np.float32 np.testing.assert_allclose(tnp, np.array(data)) - @unittest.expectedFailure + @unittest.skipIf(Device.DEFAULT=="LLVM", "broken for LLVM") def test_bf16_ones(self): - # TODO: fix this with correct bfloat16 cast t = Tensor.ones(3, 5, dtype=dtypes.bfloat16) assert t.dtype == dtypes.bfloat16 np.testing.assert_allclose(t.numpy(), np.ones((3, 5))) - @unittest.expectedFailure + @unittest.skipIf(Device.DEFAULT=="LLVM", "broken for LLVM") def test_bf16_eye(self): - # TODO: fix this with correct bfloat16 cast t = Tensor.eye(3, dtype=dtypes.bfloat16) assert t.dtype == dtypes.bfloat16 np.testing.assert_allclose(t.numpy(), np.eye(3)) @@ -256,6 +254,10 @@ class TestBitCast(unittest.TestCase): b = a.bitcast(dtypes.int32) assert b.numpy()[0] == 0x3f800000 + def test_bitcast_const(self): + a = Tensor(1, dtype=dtypes.float32).bitcast(dtypes.uint32) + np.testing.assert_equal(a.numpy(), np.array(1, dtype=np.float32).view(np.uint32)) + def test_bitcast_upcasted(self): a = Tensor.zeros(100, 4, dtype=dtypes.int32).contiguous() + 0x3f800000 b = a.bitcast(dtypes.float32) diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py index cf83136ce1..e8e0566151 100644 --- a/tinygrad/codegen/uops.py +++ b/tinygrad/codegen/uops.py @@ -95,7 +95,7 @@ class PatternMatcher: constant_folder = PatternMatcher([ # const rules ({"__name__": "root", "uop": UOps.GEP, "vin": ({"__name__": "c", "uop": UOps.CONST},)}, lambda root, c: UOp.const(root.dtype, c.arg)), - ({"__name__": "root", "uop": UOps.CAST, "vin": {"__name__": "c", "uop": UOps.CONST}}, lambda root, c: UOp.const(root.dtype, c.arg)), + ({"__name__": "root", "uop": UOps.CAST, "vin": {"__name__": "c", "uop": UOps.CONST, "arg": False}}, lambda root, c: UOp.const(root.dtype, c.arg)), # a phi without loops (len(vin)==2) is a noop ({"uop": UOps.PHI, "vin": ({}, {"__name__": "x"})}, lambda x: x), # x+-y -> x-y