Added a test case for python gradient builder. (#9207)

* Register Cos operator gradient using ORTModule's register_gradient and compare gradient against PyTorch.

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -21,6 +21,8 @@ import tempfile
 import os
 from distutils.version import LooseVersion
 
+from onnxruntime.training.ortmodule._custom_gradient_registry import register_gradient
+
 from onnxruntime.training.ortmodule import ORTModule, _utils, _io, DebugOptions, LogLevel, _fallback, _graph_execution_manager
 import _test_helpers
 
@@ -3866,3 +3868,35 @@ def test_ortmodule_determinism_flag(is_training,deterministic):
 
         from onnxruntime.training.ortmodule import _are_deterministic_algorithms_enabled
         assert _are_deterministic_algorithms_enabled() is torch.are_deterministic_algorithms_enabled()
+
+
+def test_ortmodule_gradient_builder():
+    class Model(torch.nn.Module):
+        def __init__(self):
+            super(Model, self).__init__()
+        def forward(self, x):
+            return torch.cos(x)
+
+    device = 'cuda'
+
+    @register_gradient('', 'Cos')
+    def Cos_gradient():
+        return [('Sin', ['I(0)'], ['Sin_X']),
+                ('Mul', ['Sin_X', 'GO(0)'], ['Sin_X_Times_dY']),
+                ('Neg', ['Sin_X_Times_dY'], ['GI(0)'])]
+
+    pt_model = Model().to(device)
+    ort_model = ORTModule(copy.deepcopy(pt_model))
+
+    def run_step(model, x):
+        prediction = model(x)
+        loss = prediction.sum()
+        loss.backward()
+        return prediction
+
+    pt_x = torch.randn(2, 2, device=device, requires_grad=True, dtype=torch.float32)
+    ort_x = copy.deepcopy(pt_x)
+    pt_prediction = run_step(pt_model, pt_x)
+    ort_prediction = run_step(ort_model, ort_x)
+    _test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
+    _test_helpers.assert_values_are_close(ort_x.grad, pt_x.grad )