diff --git a/orttraining/orttraining/eager/opgen/opgen/atenops.py b/orttraining/orttraining/eager/opgen/opgen/atenops.py index cd655bf08a..1e49fe86da 100644 --- a/orttraining/orttraining/eager/opgen/opgen/atenops.py +++ b/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 diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp index e3c00a9884..a24250cfd8 100644 --- a/orttraining/orttraining/eager/ort_aten.cpp +++ b/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 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 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 diff --git a/orttraining/orttraining/eager/test/ort_ops.py b/orttraining/orttraining/eager/test/ort_ops.py index e57d9c1bf1..9a9892137b 100644 --- a/orttraining/orttraining/eager/test/ort_ops.py +++ b/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).