Implement squeeze.dim for complex inputs | feat(torchlib) (#1139)

Discovered through pytorch/pytorch#113067

onnxscript/function_libs/torch_lib/ops/core.py

@@ -7261,6 +7261,15 @@ def aten_squeeze_dim(self: TTensor, dim: int) -> TTensor:
     return result
 
 
+@torch_op("aten::squeeze.dim", complex=True, trace_only=True)
+def aten_squeeze_dim_complex(self: TTensor, dim: int) -> TTensor:
+    if dim < 0:
+        # Account for the complex dimension in ONNX
+        dim = dim - 1
+
+    return aten_squeeze_dim(self, dim)
+
+
 def aten_squeeze_copy(self: TensorType) -> TensorType:
     """squeeze_copy(Tensor self) -> Tensor"""
 

onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -1373,6 +1373,15 @@ def _where_input_wrangler(
         matcher=lambda sample: not (len(sample.args) > 0 and isinstance(sample.args[0], int)),
         reason="this Aten overload only support one tensor as input and one int as args by design",
     ),
+    TorchLibOpInfo(
+        "squeeze_dim",
+        core_ops.aten_squeeze_dim_complex,
+        complex=True,
+        trace_only=True,
+    ).skip(
+        matcher=lambda sample: not (len(sample.args) > 0 and isinstance(sample.args[0], int)),
+        reason="this Aten overload only support one tensor as input and one int as args by design",
+    ),
     TorchLibOpInfo(
         "squeeze",
         core_ops.aten_squeeze,