Make doctest extension work across python versions

Signed-off-by: Jay Gu <jagu@nvidia.com>

Author: haijieg

docs/source/_ext/doctest_ext.py

@@ -125,12 +125,15 @@ class FdCaptureDocTestRunner(SphinxDocTestRunner):
     (failure reports) reaches the terminal instead of the temp file.
     """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._test_stats = {}
+
     def summarize(self, out, verbose=None):
-        groups = self._name2ft
         totalf = totalt = 0
         lines = []
-        only_default = set(groups.keys()) <= {'default'}
-        for name, (f, t) in sorted(groups.items()):
+        only_default = set(self._test_stats.keys()) <= {'default'}
+        for name, (f, t) in sorted(self._test_stats.items()):
             totalf += f
             totalt += t
             if not only_default:
@@ -156,7 +159,10 @@ def run(self, test, compileflags=None, out=None, clear_globs=True):
         orig_logger = _mod.logger
         _mod.logger = _TerminalLogger(orig_logger, saved_fd)
         try:
-            return super().run(test, compileflags, out, clear_globs)
+            result = super().run(test, compileflags, out, clear_globs)
+            f, t = self._test_stats.get(test.name, (0, 0))
+            self._test_stats[test.name] = (f + result.failed, t + result.attempted)
+            return result
         finally:
             _mod.logger = orig_logger
             self._fakeout = save_fakeout

src/cuda/tile/_stub.py

@@ -1193,22 +1193,22 @@ def _doc_atomic_rmw_op(f):
 
     Examples:
 
-        .. testcode::
-            :template: setup_only.py
+    .. testcode::
+        :template: setup_only.py
 
-            @ct.kernel
-            def kernel(x):
-                indices = ct.arange(4, dtype=ct.int32)
-                update = ct.full((4,), 10, dtype=ct.int32)
-                old = ct.{op_name}(x, indices, update)
-                print(old)
+        @ct.kernel
+        def kernel(x):
+            indices = ct.arange(4, dtype=ct.int32)
+            update = ct.full((4,), 10, dtype=ct.int32)
+            old = ct.{op_name}(x, indices, update)
+            print(old)
 
-            x = torch.ones(4, dtype=torch.int32, device='cuda')
-            ct.launch(stream, (1,), kernel, (x,))
+        x = torch.ones(4, dtype=torch.int32, device='cuda')
+        ct.launch(stream, (1,), kernel, (x,))
 
-        .. testoutput::
+    .. testoutput::
 
-            [1, 1, 1, 1]
+        [1, 1, 1, 1]
     """
 
     return f
@@ -1930,20 +1930,21 @@ def wrapped(*args, **kwargs):
     orig_doc = f.__doc__ or ""
     extra_block = _math_op_extra_block(f, indent="        ")
 
-    wrapped.__doc__ = f"""Performs {op_name} reduction on tile along the `axis`.
+    wrapped.__doc__ = f"""\
+Performs {op_name} reduction on tile along the `axis`.
 
-    Args:
-        x (Tile): input tile.
-        axis (None | const int | tuple[const int,...]): the axis for reduction.
-            The default, `axis=None`, will reduce all of the elements.
-            For `argmin` and `argmax`, tuple of axis is not supported.
-        keepdims (const bool): If true, preserves the number of dimension
-            from the input tile.{extra_block}
+Args:
+    x (Tile): input tile.
+    axis (None | const int | tuple[const int,...]): the axis for reduction.
+        The default, `axis=None`, will reduce all of the elements.
+        For `argmin` and `argmax`, tuple of axis is not supported.
+    keepdims (const bool): If true, preserves the number of dimension
+        from the input tile.{extra_block}
 
-    Returns:
-        Tile:
+Returns:
+    Tile:
 
-    """ + orig_doc
+""" + orig_doc
 
     return wrapped