TorchScript compatibility?

Hi all,

This library looks very nice 😃

Is TensorType compatible with the TorchScript compiler? As in, are the annotations transparently converted to torch.Tensor as far as torch.jit.script is concerned, allowing annotated modules/functions to be compiled? (I’m not worried about whether the type checking applied in TorchScript, just whether an annotated program that gets shape-checked in Python can be compiled down to TorchScript.)

Thanks!

[Datasciensyash]

Found another way to deal with torchscript. Just paste the code and call patch_torchscript() before exporting.

import re
import typing as tp

import torch

ttp_regexp = re.compile(r"TensorType\[[^\]]*\]")
torchtyping_replacer = "torch.Tensor"


def _replace_torchtyping(source_lines: tp.List[str]) -> tp.List[str]:

    # Join all lines
    cat_lines = "".join(source_lines)

    # Quick exit, if torchtyping is not used
    if ttp_regexp.search(cat_lines) is None:
        return source_lines

    # Replace TensorType
    cat_lines = ttp_regexp.sub(torchtyping_replacer, cat_lines)

    # Split into lines
    source_lines = cat_lines.split("\n")
    source_lines = [f"{i}\n" for i in source_lines]

    return source_lines


def _torchtyping_destruct_wrapper(func: tp.Callable) -> tp.Callable:
    def _wrap_func(obj: tp.Any, error_msg: tp.Optional[str] = None) -> tp.Tuple[tp.List[str], int, tp.Optional[str]]:
        srclines, file_lineno, filename = func(obj, error_msg)
        srclines = _replace_torchtyping(srclines)
        return srclines, file_lineno, filename

    return _wrap_func


def patch_torchscript() -> None:
    """
    Patch torchscript to work with torchtyping.

    Returns: None.

    """
    # Patch _sources if torch >= 1.10.0, else torch.jit.frontend
    if hasattr(torch, "_sources"):
        src = getattr(torch, "_sources")  # noqa: B009
    else:
        src = getattr(torch.jit, "frontend")  # noqa: B009

    src.get_source_lines_and_file = _torchtyping_destruct_wrapper(src.get_source_lines_and_file)

[patrick-kidger]

Not that I know about. As far as I know this is still a limitation in torchscript itself.

If this is a priority for you then you might like to try bringing this up with the torchscript team. They might know more about any possibilities for making this work.