printer translating thunder dtype/device to torchs

Signed-off-by: Masaki Kozuki <[email protected]>

thunder/core/prims.py

@@ -4118,8 +4118,16 @@ def printer_of_tensor_subclass_ctor(
     else:
         cls: torch._C._TensorMeta = wrapped_cls.obj
     tensors, non_tensors = arg_printables[-2:]
+    new_non_tensors = []
+    for a in non_tensors:
+        if isinstance(a, dtypes.dtype):
+            new_non_tensors.append(dtypes.to_torch_dtype(a))
+        elif isinstance(a, devices.Device):
+            new_non_tensors.append(devices.to_torch_device(a))
+        else:
+            new_non_tensors.append(a)
 
-    arg_str = ", ".join(codeutils.prettyprint(x) for x in [*tensors, *non_tensors])
+    arg_str = ", ".join(codeutils.prettyprint(x) for x in [*tensors, *new_non_tensors])
     kwarg_str = ""
 
     result_str: str

thunder/executors/torchex.py

@@ -2210,7 +2210,15 @@ def _shape_impl(t):
 
 
 def _tensor_subclass_ctor(cls, name, shape, device, dtype, requires_grad, tensors, non_tensors):
-    return cls(*tensors, *non_tensors)
+    new_non_tensors = []
+    for a in non_tensors:
+        if isinstance(a, dtypes.dtype):
+            new_non_tensors.append(to_torch_dtype(a))
+        elif isinstance(a, devices.Device):
+            new_non_tensors.append(to_torch_device(a))
+        else:
+            new_non_tensors.append(a)
+    return cls(*tensors, *new_non_tensors)
 
 
 def _bind_postprocess_of_tensor_subclass_ctor(bsym: BoundSymbol) -> None: