Showing custom types in Sphinx Docs

[a-r-j]

Hi @patrick-kidger, thanks for maintaining such a great project! :D

I'm running into some behaviour I'd like to change when rendering docs with Sphinx.

When I define a type it is rendered as a torch.Tensor in the docs e.g.:

#types.py

from torchtyping import TensorType

CoordTensor = TensorType[-1, 3]

# function.py

from .types import CoordTensor

def func(x: CoordTensor) -> CoordTensor:
    """Returns the input.

    :param x: tensor
    :type x: CoordTensor
    """"
    return x

Do you know if there is a way to get Sphinx to show the type name I have defined (and possibly make it clickable so users can jump to the type definition in the documentation?)

[a-r-j]

I managed to resolve this:

Prevents resolving to torch.Tensor

# types.py

from typing import NewType
from torchtyping import TensorType

CoordTensor = NewType("CoordTensor", TensorType[-1, 3])

And to get sphinx to render links:

# conf.py

import sphinx.ext.autodoc
sphinx.ext.autodoc.NewTypeDataDocumenter.directivetype = "class"

[patrick-kidger]

FWIW, these days I now recommmend that folks use jaxtyping, which is pretty much a straight upgrade. Despite the name (historical reasons) it also handles PyTorch, doesn't have a JAX dependency, and it appears in documentation just fine. (At least, MkDocs-based documentation. I've not tried it with Sphinx.)