Add ATenOp at::diagonal (#8838)

* Register at::diagonal for ATenOp

onnxruntime/python/tools/symbolic_shape_infer.py

@@ -187,6 +187,7 @@ class SymbolicShapeInference:
         }
         self.aten_op_dispatcher_ = {
             'aten::embedding': self._infer_Gather,
+            'aten::diagonal': self._infer_aten_diagonal,
             'aten::max_pool2d_with_indices': self._infer_aten_pool2d,
             'aten::unfold': self._infer_aten_unfold,
             'aten::argmax': self._infer_aten_argmax,
@@ -1003,6 +1004,36 @@ class SymbolicShapeInference:
                 helper.make_tensor_value_info(o, vi.type.tensor_type.elem_type,
                                               get_shape_from_sympy_shape(sympy_shape)))
 
+    def _infer_aten_diagonal(self, node):
+        sympy_shape = self._get_sympy_shape(node, 0)
+        rank = len(sympy_shape)
+        offset = self._try_get_value(node, 1)
+        dim1 = self._try_get_value(node, 2)
+        dim2 = self._try_get_value(node, 3)
+
+        assert offset is not None and dim1 is not None and dim2 is not None 
+        dim1 = handle_negative_axis(dim1, rank)
+        dim2 = handle_negative_axis(dim2, rank)
+        
+        new_shape = []
+        for dim, val in enumerate(sympy_shape):
+            if dim not in [dim1, dim2]:
+                new_shape.append(val)
+
+        shape1 = sympy_shape[dim1] 
+        shape2 = sympy_shape[dim2]
+        if offset >= 0:
+            diag_shape = sympy.Max(0, sympy.Min(shape1, shape2 - offset))
+        else:
+            diag_shape = sympy.Max(0, sympy.Min(shape1 + offset, shape2))
+        new_shape.append(diag_shape) 
+
+        if node.output[0]:
+            vi = self.known_vi_[node.output[0]]
+            vi.CopyFrom(
+                helper.make_tensor_value_info(node.output[0], self.known_vi_[node.input[0]].type.tensor_type.elem_type,
+                                              get_shape_from_sympy_shape(new_shape)))
+
     def _infer_aten_pool2d(self, node):
         sympy_shape = self._get_sympy_shape(node, 0)
         assert len(sympy_shape) == 4

orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py

@@ -86,6 +86,13 @@ def embedding_gradient():
          'GI(0)'], {'name': {'value': 'aten::embedding_backward', 'dtype': 'string'}}),
     ]
 
+@register_gradient('com.microsoft', 'ATenOp', 'aten::diagonal', '')
+def diagonal_gradient():
+    return [
+        ('Shape', ['I(0)'], ['Shape_X']),
+        (('ATenOp', 'com.microsoft'), ['GO(0)', 'Shape_X', 'I(1)', 'I(2)', 'I(3)'], [
+         'GI(0)'], {'name': {'value': 'aten::diagonal_backward', 'dtype': 'string'}}),
+    ]
 
 @register_gradient('com.microsoft', 'ATenOp', 'aten::max_pool2d_with_indices', '')
 def max_pool2d_gradient():

orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py

@@ -69,6 +69,12 @@ def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse):
     return output
 
 
+@register_symbolic('diagonal')
+def diagonal(g, self, offset, dim1, dim2):
+    return g.op("com.microsoft::ATenOp", self, offset, dim1, dim2,
+                name_s='aten::diagonal')
+
+
 @register_symbolic('max_pool2d')
 def max_pool2d(g, self, kernel_size, stride, padding, dilation, ceil_mode):
     return g.op("com.microsoft::ATenOp", self, kernel_size, stride, padding, dilation, ceil_mode,

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -782,6 +782,39 @@ def test_gradient_correctness_argmax_unfold():
         _test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
         _test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model)
 
+@pytest.mark.parametrize("offset", [-1, 0, 1])
+@pytest.mark.parametrize("dim1, dim2", ([0, 1], [0, 2], [1, 2], [2, 0]))
+def test_gradient_correctness_argmax_diagonal(offset, dim1, dim2):
+    class NeuralNetDiagonal(torch.nn.Module):
+        def __init__(self, offset=0, dim1=0, dim2=1):
+            super(NeuralNetDiagonal, self).__init__()
+            self.offset = offset
+            self.dim1 = dim1
+            self.dim2 = dim2
+
+        def forward(self, input):
+            return torch.diagonal(input, self.offset, self.dim1, self.dim2)
+
+    N, D, H = 16, 256, 128
+    device = 'cuda'
+    pt_model = NeuralNetDiagonal(offset, dim1, dim2).to(device)
+    ort_model = ORTModule(copy.deepcopy(pt_model))
+
+    def run_step(model, input):
+        prediction = model(input)
+        loss = prediction.sum()
+        loss.backward()
+        return prediction
+
+    for _ in range(10):
+        pt_input = torch.rand((N, D, H), device=device, requires_grad=True)
+        ort_input = copy.deepcopy(pt_input)
+        pt_prediction = run_step(pt_model, pt_input)
+        ort_prediction = run_step(ort_model, ort_input)
+
+        _test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
+        _test_helpers.assert_values_are_close(ort_input.grad, pt_input.grad)
+
 def test_module_with_non_differential_output():
     device = 'cuda'
     N, D_in, H, D_out = 32, 128, 64, 10