feat(atenlib): add ops(squeeze) (#460)

Author: xiaowuhu

onnxscript/function_libs/torch_aten/ops/core.py

@@ -5053,10 +5053,20 @@ def aten_square(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_squeeze(self: TensorType) -> TensorType:
+@torch_op("aten::squeeze", trace_only=True)
+def aten_squeeze(self: TTensor, dim: Optional[int] = None) -> TTensor:
     """squeeze(Tensor(a) self) -> Tensor(a)"""
 
-    raise NotImplementedError()
+    if op.OptionalHasElement(dim):
+        rank = op.Size(op.Shape(self))
+        if rank == 0:
+            self = op.Reshape(self, op.Constant(value_ints=[-1]))
+        dims = op.Reshape(dim, op.Constant(value_ints=[-1]))
+        result = op.Squeeze(self, dims)
+    else:
+        result = op.Squeeze(self)
+
+    return result
 
 
 def aten_squeeze_copy(self: TensorType) -> TensorType:

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -412,6 +412,7 @@ def _where_input_wrangler(
     ),
     "ones_like": core_ops.aten_ones_like,
     "slice": core_ops.aten_slice,
+    "squeeze": core_ops.aten_squeeze,
     "sum": (core_ops.aten_sum_dim_IntList, _sum_input_wrangler),
     "transpose": core_ops.aten_transpose,
     "zeros_like": core_ops.aten_zeros_like,
@@ -556,6 +557,13 @@ def _where_input_wrangler(
         matcher=lambda sample: len(sample.args[0]) == 0,
         reason="Empty perm is not supported",
     ),
+    skip(
+        "squeeze",
+        matcher=lambda sample: len(sample.args) > 0
+        and len(sample.input.shape) > 0
+        and sample.input.shape[sample.args[0]] != 1,
+        reason="Cannot select an axis to squeeze out which has size not equal to one",
+    ),
 )
 
 duplicate_opinfo(