Enables reference checking during doc generation

Adds `-n` to `sphinx-build` which enables checks for cross-document references.
This will allow us to ensure that links that point to other parts of the docs
are always valid.

.github/workflows/tripy-l0.yml

@@ -61,7 +61,7 @@ jobs:
         options: --gpus all -v ${{ github.workspace }}/tripy:/tripy
         run: |
           python3 docs/generate_rsts.py
-          sphinx-build build/doc_sources build/docs -c docs/ -j 4 -W
+          sphinx-build build/doc_sources build/docs -c docs/ -j 4 -W -n
 
     - name: run-test
       uses: addnab/docker-run-action@v3

.github/workflows/tripy-l1.yml

@@ -40,7 +40,7 @@ jobs:
       run: |
         cd /tripy/
         python3 docs/generate_rsts.py
-        sphinx-build build/doc_sources build/docs -c docs/ -j 4 -W
+        sphinx-build build/doc_sources build/docs -c docs/ -j 4 -W -n
         cp docs/packages.html build/docs/
 
     - uses: actions/upload-pages-artifact@v3

tripy/docs/README.md

@@ -7,7 +7,7 @@ This directory includes all the source files for the public API documentation.
 You can build the documentation locally in the development container by running:
 ```bash
 python3 docs/generate_rsts.py
-sphinx-build build/doc_sources build/docs -c docs/ -j 4 -W
+sphinx-build build/doc_sources build/docs -c docs/ -j 4 -W -n
 ```
 To view the documentation, you can open `build/docs/index.html` in a browser.
 

tripy/docs/conf.py

@@ -51,6 +51,18 @@
 autodoc_typehints_format = "short"
 python_use_unqualified_type_names = True
 
+nitpick_ignore = {
+    ("py:class", "tripy.types.ShapeLike"),
+    ("py:class", "tripy.types.TensorLike"),
+    ("py:class", "tripy.types.NestedNumberSequence"),
+    ("py:class", "Tensor"),
+    ("py:class", "Shape"),
+}
+nitpick_ignore_regex = {
+    ("py:class", r"numbers\.Number"),
+    ("py:class", r"collections\..*"),
+}
+
 autodoc_default_options = {
     "members": True,
     "no-undoc-members": True,

tripy/tripy/backend/api/executable.py

@@ -162,9 +162,7 @@ def add(a, b):
                             )
             raise
 
-        from tripy.utils.stack_info import StackInfo
-
-        output_tensors = [Tensor(output, stack_info=StackInfo([])) for output in executor_outputs]
+        output_tensors = [Tensor(output, fetch_stack_info=False) for output in executor_outputs]
         if len(output_tensors) == 1:
             output_tensors = output_tensors[0]
         return output_tensors

tripy/tripy/frontend/module/convolution.py → tripy/tripy/frontend/module/conv.py

@@ -83,7 +83,7 @@ def __init__(
         self.dtype = dtype
 
 
-@export.public_api(document_under="operations/modules")
+@export.public_api(document_under="operations/modules", autodoc_options=[":no-show-inheritance:"])
 @dataclass
 class Conv(ConvBase):
     r"""

tripy/tripy/frontend/module/conv_transpose.py

@@ -21,11 +21,11 @@
 
 from tripy import export
 from tripy.common import datatype
-from tripy.frontend.module.convolution import ConvBase
+from tripy.frontend.module.conv import ConvBase
 from tripy.frontend.module.parameter import DefaultParameter, Parameter
 
 
-@export.public_api(document_under="operations/modules")
+@export.public_api(document_under="operations/modules", autodoc_options=[":no-show-inheritance:"])
 @dataclass
 class ConvTranspose(ConvBase):
     r"""

tripy/tripy/frontend/shape.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 #
 
-from typing import Optional, Sequence, Union
+from typing import Any, Optional, Sequence, Union
 
 from tripy import export, utils
 from tripy.common.datatype import int32
@@ -41,7 +41,7 @@ class Shape(Tensor):
 
     def __init__(
         self,
-        data: Union[Sequence, Tensor, "np.ndarray", "cp.ndarray", "torch.Tensor", "jnp.ndarray"],
+        data: Any,
         name: Optional[str] = None,
     ) -> None:
         r"""
@@ -297,7 +297,7 @@ class ShapeScalar(Tensor):
 
     def __init__(
         self,
-        data: Union[Sequence, Tensor, "np.ndarray", "cp.ndarray", "torch.Tensor", "jnp.ndarray"],
+        data: Any,
         name: Optional[str] = None,
     ) -> None:
         r"""

tripy/tripy/frontend/tensor.py

@@ -16,7 +16,9 @@
 #
 
 from textwrap import indent
-from typing import Any, List, Optional, Sequence, Union
+from typing import Any, Optional
+
+import mlir_tensorrt.runtime.api as runtime
 
 # Import ops to populate the registry before we define our Tensor class
 import tripy.frontend.ops
@@ -27,8 +29,7 @@
 from tripy.common.exception import raise_error
 from tripy.frontend.ops.registry import TENSOR_METHOD_REGISTRY
 from tripy.frontend.trace.ops import Storage
-
-import mlir_tensorrt.runtime.api as runtime
+from tripy.utils.stack_info import StackInfo
 
 
 class TensorMeta(type):
@@ -73,19 +74,21 @@ def _get_unique_name(cls):
 
     def __init__(
         self,
-        data: Union[List, "np.ndarray", "cp.ndarray", "torch.Tensor", "jnp.ndarray"],
+        data: Any,
         dtype: Optional["tripy.dtype"] = None,
         device: Optional["tripy.device"] = None,
         name: Optional[str] = None,
-        stack_info: Optional["StackInfo"] = None,
+        fetch_stack_info: bool = True,
     ) -> None:
         """
         Args:
             data: The data with which to initialize the tensor.
             dtype: The data type of the tensor.
             device: The device on which to allocate the tensor.
             name: The name of the tensor. If provided, this must be a unique string.
-            stack_info: The stack infomation of the tensor.
+            fetch_stack_info: Whether to fetch stack information for the tensor.
+                Stack information allows Tripy to generate much higher quality error
+                messages at the cost of a small overhead when initializing the tensor.
 
         .. code-block:: python
             :linenos:
@@ -95,13 +98,12 @@ def __init__(
         """
         from tripy.frontend.trace.tensor import TraceTensor
 
-        # We include code for everything above the `BaseTraceOp.build` function, which is called at most
-        # this many stack frames above the constructor.
-        STACK_DEPTH_OF_BUILD = 4
-        # not using utils.default() because it always evaluates the `default` argument.
-        stack_info = (
-            stack_info if stack_info is not None else utils.get_stack_info(include_code_index=STACK_DEPTH_OF_BUILD)
-        )
+        stack_info = StackInfo([])
+        if fetch_stack_info:
+            # We include code for everything above the `BaseTraceOp.build` function, which is called at most
+            # this many stack frames above the constructor.
+            STACK_DEPTH_OF_BUILD = 4
+            stack_info = utils.get_stack_info(include_code_index=STACK_DEPTH_OF_BUILD)
 
         name = name if name is not None else Tensor._get_unique_name()
 

tripy/tripy/types.py

@@ -21,40 +21,36 @@
 
 import numbers
 import sys
-from typing import Union, Sequence
+from typing import Sequence, Union
 
 from tripy import export
 
 export.public_api(autodoc_options=[":no-members:", ":no-special-members:"])(sys.modules[__name__])
 
 NestedNumberSequence = export.public_api(
     document_under="types.rst",
-    autodoc_options=[":no-index:"],
     module=sys.modules[__name__],
     symbol="NestedNumberSequence",
     doc="""
         Denotes the recursive type annotation for sequences of Python numbers, possibly nested to an arbitrary depth.
-        Tripy often automatically converts these sequences to `tp.Tensor`.
+        Tripy often automatically converts these sequences to :class:`Tensor`.
         """,
 )(Union[numbers.Number, Sequence["tripy.types.NestedNumberSequence"]])
 
 TensorLike = export.public_api(
     document_under="types.rst",
-    autodoc_options=[":no-index:"],
     module=sys.modules[__name__],
     symbol="TensorLike",
     doc="""
-        Type annotation for a parameter that is either a Tripy `Tensor` or a Python sequence that can be automatically converted into one.
+        Type annotation for a parameter that is either a Tripy :class:`Tensor` or a Python sequence that can be automatically converted into one.
         """,
 )(Union["tripy.Tensor", "tripy.types.NestedNumberSequence"])
 
-
 ShapeLike = export.public_api(
     document_under="types.rst",
-    autodoc_options=[":no-index:"],
     module=sys.modules[__name__],
     symbol="ShapeLike",
     doc="""
-        Type annotation for a parameter that is either a Tripy `Shape` or Python sequence that can be automatically converted into one.
+        Type annotation for a parameter that is either a Tripy :class:`Shape` or Python sequence that can be automatically converted into one.
         """,
 )(Union["tripy.Shape", Sequence[Union[int, "tripy.ShapeScalar"]]])