Add is_cpu op (#2758)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Masato Shinokawa <[email protected]>

Author: 3 people

thunder/tests/opinfos.py

@@ -555,6 +555,20 @@ def _is_cuda_torch(x: torch.Tensor) -> bool:
 tensor_properties.append(is_cuda_opinfo)
 
 
+def _is_cpu_torch(x: torch.Tensor) -> bool:
+    return x.is_cpu
+
+
+is_cpu_opinfo = OpInfo(
+    ltorch.is_cpu,
+    sample_input_generator=elementwise_unary_generator,
+    torch_reference=_is_cpu_torch,
+    dtypes=(datatypes.all_dtypes),
+)
+
+tensor_properties.append(is_cpu_opinfo)
+
+
 def _is_nested_torch(x: torch.Tensor) -> bool:
     return x.is_nested
 

thunder/torch/__init__.py

@@ -333,6 +333,14 @@ def is_cuda(a: TensorLike, /) -> bool:
     return a.device.devicetype is devices.DeviceType.CUDA
 
 
+@torchsymbol(torch.Tensor.is_cpu, is_property=True, id="torch.Tensor.is_cpu")
+def is_cpu(a: TensorLike, /) -> bool:
+    return a.device.devicetype is devices.DeviceType.CPU
+
+
+register_method("is_cpu", is_cpu)
+
+
 @torchsymbol(torch.polar, id="torch.polar")
 def polar(abs_: TensorLike, angle: TensorLike):
     # Decompose polar into real/imag using Thunder ops and construct complex via mapped torch.complex