Eager mode - argmax_out: set output tensor (#12233)

This change updates the implementation or te argmax_out operator to 1)
set the output tensor correctly and 2) remove the unnecessary use of a
temporary tensor to store intermediate result of onnx ArgMax operation.

Previously, the argmax_out operator did not correctly update the out
tensor - it replaced the OrtValue instead of the memory backing the
OrtValue . To properly update the output tensor, we need to calculate
the expected shape of the out tensor.

We add the helper function calculate_reduction_shape to calculate the
shape of the reduced tensor from the input tensor, dimension to reduce,
and option to keep the reduced dimension or not. This is based on the
utility functions in aten/src/ATen/native/ReduceOpsUtils.h in the
PyTorch repository, but is tailored to be a bit more specific to our
current needs.

Notes:

We considered just directly leveraging PyTorch's utility functions (e.g.
get_reduction_shape) to calculate the shape of the reduced tensor from
aten/src/ATen/native/ReduceOpsUtils.h in the PyTorch repository, but
including this header file resulted in warnings around unused functions
that we need to handle. As we only need a limited functionality at the
moment, we instead implemented our own utility function to calculate the
reduction shape for our specific current needs. If we need a utility
function to more generally calculate the reduction shape, we could
consider switching to leveraging the utility methods in PyTorch.

orttraining/orttraining/eager/ort_aten.cpp

@@ -365,6 +365,37 @@ OrtValue CastToType(onnxruntime::ORTInvoker& invoker, const OrtValue& input, at:
   return output[0];
 }
 
+/*
+ * Utility method to calculate the resulting shape of tensor after a reduction operation.
+ *
+ * @param dimToReduce The dimension to reduce. If null, then shape is 0 dimension.
+ * @param keepdim Whether to retain dim or not. Ignored if dimToReduce is null.
+ */
+inline at::DimVector calculate_reduction_shape(
+  const at::Tensor& self,
+  c10::optional<int64_t> dimToReduce,
+  bool keepdim) {
+  at::DimVector shape;
+
+  // If we have dim value, then reduce that dimension.
+  // else, return empty shape (corresponding to 0-D tensor)
+  if (dimToReduce.has_value()) {
+    shape = at::DimVector(self.sizes());
+    int64_t effectiveDimToReduce = *dimToReduce;
+    at::maybe_wrap_dims_n(&effectiveDimToReduce, 1, self.dim());
+
+    if (keepdim) {
+      shape[effectiveDimToReduce] = 1;
+    } else {
+      shape.erase(shape.begin() + effectiveDimToReduce);
+    }
+  } else {
+    shape = at::DimVector();
+  }
+
+  return shape;
+}
+
 /*
  * Utility function for resizing output tensor
  * Only resizes if:
@@ -807,17 +838,27 @@ at::Tensor& out) {
   auto ort_input_self =
     create_ort_value(invoker, dim.has_value() ? self : self.reshape({-1}));
 
-  // Remove this hand signature once the generator can support this one line below.
   int64_t l_axis = dim.has_value() ? *dim : 0;
+  bool keepdim_effective_value = dim.has_value() ? keepdim : false;
 
   NodeAttributes attrs(2);
   attrs["axis"] = create_ort_attribute(
-  "axis", l_axis, at::ScalarType::Int);
+    "axis", l_axis, at::ScalarType::Int);
   attrs["keepdims"] = create_ort_attribute(
-  "keepdims", keepdim, at::ScalarType::Int);
+    "keepdims", keepdim_effective_value, at::ScalarType::Bool);
 
   std::vector<OrtValue> ort_outputs_0_ArgMax(1);
 
+  // Calculate the size of the out tensor, based on self tensor, dimension input, and keepdim input
+  auto shape = calculate_reduction_shape(self, dim, keepdim);
+
+  resize_output(invoker,
+                dynamic_cast<ORTTensorImpl*>(out.unsafeGetTensorImpl()),
+                at::IntArrayRef{shape});
+
+  auto ort_input_out = create_ort_value(invoker, out);
+  ort_outputs_0_ArgMax[0] = ort_input_out;
+
   auto status = invoker.Invoke("ArgMax", {
   std::move(ort_input_self),
   }, ort_outputs_0_ArgMax, &attrs);
@@ -826,12 +867,6 @@ at::Tensor& out) {
   throw std::runtime_error(
   "ORT return failure status:" + status.ErrorMessage());
 
-  at::TensorOptions tensor_options = out.options();
-
-  // generator also needs to do this to handle the out param!
-  out = aten_tensor_from_ort(
-  std::move(ort_outputs_0_ArgMax[0]),
-  tensor_options);
   return out;
 }
 

orttraining/orttraining/eager/test/ort_ops.py

@@ -258,13 +258,63 @@ class OrtOpTests(unittest.TestCase):
 
     def test_argmax(self):
         device = self.get_device()
-        cpu_tensor = torch.rand(3, 5)
+        cpu_tensor = torch.rand(3, 5, 7, 8)
         ort_tensor = cpu_tensor.to(device)
+
+        # Scenario: basic (no dim parameters)
+        cpu_result = torch.argmax(cpu_tensor)
+        ort_result = torch.argmax(ort_tensor)
+        assert torch.allclose(cpu_result, ort_result.cpu())
+        assert cpu_result.dim() == ort_result.dim()
+
+        # Scenario: specify dim parameter
         cpu_result = torch.argmax(cpu_tensor, dim=1)
         ort_result = torch.argmax(ort_tensor, dim=1)
         assert torch.allclose(cpu_result, ort_result.cpu())
         assert cpu_result.dim() == ort_result.dim()
 
+        # Scenario: specify dim and keepdim parameters
+        cpu_result = torch.argmax(cpu_tensor, dim=1, keepdim=True)
+        ort_result = torch.argmax(ort_tensor, dim=1, keepdim=True)
+        assert torch.allclose(cpu_result, ort_result.cpu())
+        assert cpu_result.dim() == ort_result.dim()
+
+        # Scenario: specify negative dim value
+        cpu_result = torch.argmax(cpu_tensor, dim=-1)
+        ort_result = torch.argmax(ort_tensor, dim=-1)
+        assert torch.allclose(cpu_result, ort_result.cpu())
+        assert cpu_result.dim() == ort_result.dim()
+
+        # Scenario: basic out (no dim parameters)
+        cpu_out_tensor = torch.tensor([], dtype=torch.long)
+        ort_out_tensor = cpu_out_tensor.to(device)
+        cpu_result = torch.argmax(cpu_tensor, out=cpu_out_tensor)
+        ort_result = torch.argmax(ort_tensor, out=ort_out_tensor)
+        assert torch.allclose(cpu_result, ort_result.cpu())
+        assert cpu_result.dim() == ort_result.dim()
+        assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu())
+        assert cpu_out_tensor.dim() == ort_out_tensor.dim()
+
+        # Scenario: out with dim parameter
+        cpu_out_tensor = torch.tensor([], dtype=torch.long)
+        ort_out_tensor = cpu_out_tensor.to(device)
+        cpu_result = torch.argmax(cpu_tensor, dim=1, out=cpu_out_tensor)
+        ort_result = torch.argmax(ort_tensor, dim=1, out=ort_out_tensor)
+        assert torch.allclose(cpu_result, ort_result.cpu())
+        assert cpu_result.dim() == ort_result.dim()
+        assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu())
+        assert cpu_out_tensor.dim() == ort_out_tensor.dim()
+
+        # Scenario: out with dim and keepdim parameters
+        cpu_out_tensor = torch.tensor([], dtype=torch.long)
+        ort_out_tensor = cpu_out_tensor.to(device)
+        cpu_result = torch.argmax(cpu_tensor, dim=1, keepdim=True, out=cpu_out_tensor)
+        ort_result = torch.argmax(ort_tensor, dim=1, keepdim=True, out=ort_out_tensor)
+        assert torch.allclose(cpu_result, ort_result.cpu())
+        assert cpu_result.dim() == ort_result.dim()
+        assert torch.allclose(cpu_out_tensor, ort_out_tensor.cpu())
+        assert cpu_out_tensor.dim() == ort_out_tensor.dim()
+
     def test_masked_select(self):
         device = self.get_device()
         cpu_tensor = torch.randn(3, 4)