Lironkesem/unsqueeze_and_squeeze (#12421)

orttraining/orttraining/eager/opgen/opgen/atenops.py

@@ -178,6 +178,9 @@ hand_implemented = {
     "aten::nll_loss_forward.output": MakeTorchFallback(),
     "aten::nll_loss_backward.grad_input": MakeTorchFallback(),
     "aten::_log_softmax_backward_data.out": MakeTorchFallback(),
+    "aten::squeeze.dim": Squeeze("self", "dim"),
+    "aten::squeeze": SignatureOnly(),
+    "aten::unsqueeze": Unsqueeze(data="self", axes="dim"),
 }
 
 # If the aten op expects a specific output type that differs from self

orttraining/orttraining/eager/ort_aten.cpp

@@ -1121,6 +1121,36 @@ at::Tensor& mm_out(
   return out;
 }
 
+// aten::squeeze(Tensor(a) self) -> Tensor(a)
+at::Tensor squeeze(
+  const at::Tensor& self) {
+  ORT_LOG_FN(self);
+
+  if (
+       std::vector<at::ScalarType> supportedTypes =
+          {at::kBFloat16, at::kBool, at::kByte, at::kDouble, at::kFloat, at::kHalf, at::kInt, at::kLong, at::kShort};
+      !IsSupportedType(self, supportedTypes))
+    return at::native::call_fallback_fn<
+      &at::native::cpu_fallback,
+      ATEN_OP(squeeze)>::call(self);
+
+  auto& invoker = GetORTInvoker(self.device());
+
+  auto ort_input_0_self = create_ort_value(invoker, self);
+
+  std::vector<OrtValue> ort_outputs_0_Squeeze(1);
+
+  auto status = invoker.Invoke("Squeeze", {
+    std::move(ort_input_0_self),
+  }, ort_outputs_0_Squeeze, nullptr);
+  CHECK_STATUS(status);
+
+  at::TensorOptions tensor_options = self.options();
+  return aten_tensor_from_ort(
+    std::move(ort_outputs_0_Squeeze[0]),
+    tensor_options);
+}
+
 }  // namespace aten
 
 // #pragma endregion

orttraining/orttraining/eager/test/ort_ops.py

@@ -444,6 +444,33 @@ class OrtOpTests(unittest.TestCase):
         with self.assertRaises(RuntimeError):
             torch.mm(ort_mat1, ort_not_matrix)
 
+    def test_squeeze(self):
+        device = self.get_device()
+        cpu_tensor = torch.zeros(2, 1, 2, 1, 2)
+        ort_tensor = cpu_tensor.to(device)
+
+        cpu_result1 = torch.squeeze(cpu_tensor)
+        ort_result1 = torch.squeeze(ort_tensor)
+
+        cpu_result2 = torch.squeeze(cpu_tensor, 1)
+        ort_result2 = torch.squeeze(ort_tensor, 1)
+
+        assert torch.equal(cpu_result1, ort_result1.cpu())
+        assert torch.equal(cpu_result2, ort_result2.cpu())
+
+    def test_unsqueeze(self):
+        device = self.get_device()
+        cpu_tensor = torch.tensor([1, 2, 3, 4])
+        ort_tensor = cpu_tensor.to(device)
+
+        cpu_result1 = torch.unsqueeze(cpu_tensor, 0)
+        ort_result1 = torch.unsqueeze(ort_tensor, 0)
+        cpu_result2 = torch.unsqueeze(cpu_tensor, 1)
+        ort_result2 = torch.unsqueeze(ort_tensor, 1)
+
+        assert torch.equal(cpu_result1, ort_result1.cpu())
+        assert torch.equal(cpu_result2, ort_result2.cpu())
+
     ################################ parameterized test follow #######################################
     # OPS - is a list of [test_operator, tested_tensor=torch.rand (6)].
     # The default value for tested_tensor is torch.rand (6)- size of 6 uniform distribution on the interval [0, 1).