feat: added tensor long (#402)

thunder/tests/opinfos.py

@@ -2919,6 +2919,28 @@ def type_as_sample_generator(op, device, dtype, requires_grad, **kwargs):
 data_movement_ops.append(type_as_sample)
 
 
+def long_sample_generator(op, device, dtype, requires_grad, **kwargs):
+    make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
+
+    # input shape
+    shapes = (
+        (),
+        (0,),
+        (2,),
+        (1, 2),
+        (1, 2, 3),
+    )
+
+    for shape in shapes:
+        yield SampleInput(make(shape))
+
+
+long_opinfo = OpInfo(
+    ltorch.long,
+    sample_input_generator=long_sample_generator,
+    torch_reference=torch.Tensor.long,
+)
+
 opinfos.extend(data_movement_ops)
 
 #

thunder/torch/__init__.py

@@ -435,6 +435,11 @@ def type_as(a: TensorProxy, b: TensorProxy, /) -> TensorProxy:
     return to(a, b.true_dtype)
 
 
+@torchsymbol(torch.Tensor.long, is_method=True)
+def long(a: TensorLike, /, memory_format: torch.memory_format = torch.preserve_format) -> TensorLike:
+    return to(a, dtype=dtypes.int64, memory_format=memory_format)
+
+
 #
 # Tensor creation operations
 #