Add uint64 to thunder->torch dtype map (#2519)

Author: crcrpar

thunder/core/dtypes.py

@@ -239,7 +239,7 @@ def __init__(self, name, shortname, *, bytes, is_weak):
 
 boolean_dtypes = {bool8, bool8_, bool}
 
-integer_dtypes = {d for d in all_dtypes if isinstance(d, exact)} | {bool, int}
+integer_dtypes = {d for d in all_dtypes if isinstance(d, exact)} | {bool, int} | {uint64, uint64_}
 
 nonboolean_integer_dtypes = {d for d in integer_dtypes if (not isinstance(d, bool_) and d is not bool)}
 
@@ -264,7 +264,7 @@ def __init__(self, name, shortname, *, bytes, is_weak):
 
 weak_dtypes = {d for d in all_dtypes if d.is_weak} | all_numbertypes
 
-strong_dtypes = {d for d in all_dtypes if not d.is_weak}
+strong_dtypes = {d for d in all_dtypes if not d.is_weak} | {uint64}
 
 
 def is_weak_dtype(dtype):
@@ -384,9 +384,11 @@ def corresponding_complex_dtype(dtype):
 
 
 _name_to_dtype_map = {dtype.full_name: dtype for dtype in all_dtypes}
+_name_to_dtype_map.update({uint64.full_name: uint64, uint64_.full_name: uint64_})
 _strong_dtype_to_weak_dtype_map = {
     dtype: _name_to_dtype_map[f"{dtype.full_name}_"] for dtype in all_dtypes if not dtype.is_weak
 }
+_strong_dtype_to_weak_dtype_map.update({uint64: uint64_})
 
 _weak_dtype_to_strong_dtype_map = {v: k for k, v in _strong_dtype_to_weak_dtype_map.items()}
 _weak_dtype_to_strong_dtype_map.update(
@@ -395,6 +397,7 @@ def corresponding_complex_dtype(dtype):
         int: int64,
         float: float32,
         complex: complex64,
+        uint64_: uint64,
     }
 )
 
@@ -524,6 +527,12 @@ def are_same_dtypes(a, b, *, weak_and_strong_are_equivalent=True):
         if hasattr(torch, dtype.full_name.rstrip("_"))
     }
 )
+_thunder_to_torch_dtype_map.update(
+    {
+        uint64: torch.uint64,
+        uint64_: torch.uint64,
+    }
+)
 
 _torch_to_thunder_dtype_map = {
     v: k