add unit tests for to_dtype (tinygrad#8217)

* add unit test for to_dtype * add unit test for to_dtype --------- Co-authored-by: pkotzbach <[email protected]>
chenyuxyz · Dec 13, 2024 · c1b79c1 · c1b79c1
1 parent 8a50868
commit c1b79c1
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/test/test_dtype.py b/test/test_dtype.py
@@ -4,7 +4,7 @@
 from typing import Any, List
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, DEBUG, CI
-from tinygrad.dtype import DType, DTYPES_DICT, ImageDType, PtrDType, least_upper_float, least_upper_dtype, truncate_fp16
+from tinygrad.dtype import DType, DTYPES_DICT, ImageDType, PtrDType, least_upper_float, least_upper_dtype, truncate_fp16, to_dtype
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
 from hypothesis import given, settings, strategies as strat
@@ -854,5 +854,18 @@ def test_max_w_alu(self):
         t = Tensor([[1, 2], [3, 4]], dtype=d)
         (t*t).max().item()
 
+class TestToDtype(unittest.TestCase):
+  def test_dtype_to_dtype(self):
+    dtype = dtypes.int32
+    res = to_dtype(dtype)
+    self.assertIsInstance(res, DType)
+    self.assertEqual(res, dtypes.int32)
+
+  def test_str_to_dtype(self):
+    dtype = "int32"
+    res = to_dtype(dtype)
+    self.assertIsInstance(res, DType)
+    self.assertEqual(res, dtypes.int32)
+
 if __name__ == '__main__':
   unittest.main()