Support torch.square natively (#2329)

Author: crcrpar

thunder/tests/opinfos.py

@@ -2192,6 +2192,14 @@ def imag_error_generator(op, device, **kwargs):
 elementwise_unary_ops.append(clone_opinfo)
 
 
+square_opinfo = OpInfo(
+    ltorch.square,
+    sample_input_generator=elementwise_unary_generator,
+    torch_reference=_elementwise_unary_torch(torch.square),
+)
+elementwise_unary_ops.append(square_opinfo)
+
+
 # Puts all opinfos into the "opinfos" list
 opinfos.extend(elementwise_unary_ops)
 

thunder/torch/__init__.py

@@ -2248,6 +2248,25 @@ def threshold_(a: TensorProxy, /, threshold: float, value: float) -> TensorLike:
 _inplace_to_out_of_place[threshold_] = threshold, -1
 
 
+@torchsymbol(torch.square, is_method=True)
+def square(a):
+    if isinstance(dtypes.to_dtype(a), dtypes.bool_):
+        a = clang.maybe_convert_to_dtype(a, dtypes.int64)
+    return a * a
+
+
+@torchsymbol(torch.square_, is_method=True, tags=(prims.OpTags.IN_PLACE,))
+def square_(a):
+    utils.check(
+        not isinstance(dtypes.to_dtype(a), dtypes.bool_),
+        lambda: f"Result type of {dtypes.int64} cannot be stored into {dtypes.to_dtype(a)}",
+    )
+    return _copy_(a, square(a))
+
+
+_inplace_to_out_of_place[square_] = square, -1
+
+
 #
 # Elementwise binary operations
 #

thunder/torch/default_torch_ops.py

@@ -252,7 +252,6 @@
         torch.split_with_sizes,
         torch.split_with_sizes_copy,
         torch.spmm,
-        torch.square,
         torch.squeeze_copy,
         torch.sspaddmm,
         torch.std_mean,
@@ -548,7 +547,6 @@
         torch.Tensor.sparse_dim,
         torch.Tensor.sparse_mask,
         torch.Tensor.split_with_sizes,
-        torch.Tensor.square,
         torch.Tensor.sspaddmm,
         torch.Tensor.stft,
         torch.Tensor.subtract,