Add files using upload-large-folder tool

0e9bdcb verified over 1 year ago

17.7 kB

import ast
import builtins
import dis
import enum
import inspect
import re
import warnings

from textwrap import dedent
from typing import Type

import torch

from torch._C import (
    AnyType,
    AwaitType,
    BoolType,
    ComplexType,
    DeviceObjType,
    DictType,
    EnumType,
    FloatType,
    FutureType,
    InterfaceType,
    IntType,
    ListType,
    NoneType,
    NumberType,
    OptionalType,
    StreamObjType,
    StringType,
    TensorType,
    TupleType,
    UnionType,
)
from torch._sources import get_source_lines_and_file
from .._jit_internal import (  # type: ignore[attr-defined]
    _Await,
    _qualified_name,
    Any,
    BroadcastingList1,
    BroadcastingList2,
    BroadcastingList3,
    Dict,
    Future,
    is_await,
    is_dict,
    is_future,
    is_ignored_fn,
    is_list,
    is_optional,
    is_tuple,
    is_union,
    List,
    Optional,
    Tuple,
    Union,
)
from ._state import _get_script_class

if torch.distributed.rpc.is_available():
    from torch._C import RRefType
    from .._jit_internal import is_rref, RRef

from torch._ops import OpOverloadPacket


class Module:
    def __init__(self, name, members):
        self.name = name
        self.members = members

    def __getattr__(self, name):
        try:
            return self.members[name]
        except KeyError:
            raise RuntimeError(
                f"Module {self.name} has no member called {name}"
            ) from None


class EvalEnv:
    env = {
        "torch": Module("torch", {"Tensor": torch.Tensor}),
        "Tensor": torch.Tensor,
        "typing": Module("typing", {"Tuple": Tuple}),
        "Tuple": Tuple,
        "List": List,
        "Dict": Dict,
        "Optional": Optional,
        "Union": Union,
        "Future": Future,
        "Await": _Await,
    }

    def __init__(self, rcb):
        self.rcb = rcb
        if torch.distributed.rpc.is_available():
            self.env["RRef"] = RRef

    def __getitem__(self, name):
        if name in self.env:
            return self.env[name]
        if self.rcb is not None:
            return self.rcb(name)
        return getattr(builtins, name, None)


def get_signature(fn, rcb, loc, is_method):
    if isinstance(fn, OpOverloadPacket):
        signature = try_real_annotations(fn.op, loc)
    else:
        signature = try_real_annotations(fn, loc)
    if signature is not None and is_method:
        # If this is a method, then the signature will include a type for
        # `self`, but type comments do not contain a `self`. So strip it
        # away here so everything is consistent (`inspect.ismethod` does
        # not work here since `fn` is unbound at this point)
        param_types, return_type = signature
        param_types = param_types[1:]
        signature = (param_types, return_type)

    if signature is None:
        type_line, source = None, None
        try:
            source = dedent("".join(get_source_lines_and_file(fn)[0]))
            type_line = get_type_line(source)
        except TypeError:
            pass
        # This might happen both because we failed to get the source of fn, or
        # because it didn't have any annotations.
        if type_line is not None:
            signature = parse_type_line(type_line, rcb, loc)

    return signature


def is_function_or_method(the_callable):
    # A stricter version of `inspect.isroutine` that does not pass for built-in
    # functions
    return inspect.isfunction(the_callable) or inspect.ismethod(the_callable)


def is_vararg(the_callable):
    if not is_function_or_method(the_callable) and callable(the_callable):  # noqa: B004
        # If `the_callable` is a class, de-sugar the call so we can still get
        # the signature
        the_callable = the_callable.__call__

    if is_function_or_method(the_callable):
        return inspect.getfullargspec(the_callable).varargs is not None
    else:
        return False


def get_param_names(fn, n_args):
    if isinstance(fn, OpOverloadPacket):
        fn = fn.op

    if (
        not is_function_or_method(fn)
        and callable(fn)
        and is_function_or_method(fn.__call__)
    ):  # noqa: B004
        # De-sugar calls to classes
        fn = fn.__call__

    if is_function_or_method(fn):
        if is_ignored_fn(fn):
            fn = inspect.unwrap(fn)
        return inspect.getfullargspec(fn).args
    else:
        # The `fn` was not a method or function (maybe a class with a __call__
        # method, so use a default param name list)
        return [str(i) for i in range(n_args)]


def check_fn(fn, loc):
    # Make sure the function definition is not a class instantiation
    try:
        source = dedent("".join(get_source_lines_and_file(fn)[0]))
    except (OSError, TypeError):
        return
    if source is None:
        return

    py_ast = ast.parse(source)
    if len(py_ast.body) == 1 and isinstance(py_ast.body[0], ast.ClassDef):
        raise torch.jit.frontend.FrontendError(
            loc,
            f"Cannot instantiate class '{py_ast.body[0].name}' in a script function",
        )
    if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef):
        raise torch.jit.frontend.FrontendError(
            loc, "Expected a single top-level function"
        )


def _eval_no_call(stmt, glob, loc):
    """Evaluate statement as long as it does not contain any method/function calls"""
    bytecode = compile(stmt, "", mode="eval")
    for insn in dis.get_instructions(bytecode):
        if "CALL" in insn.opname:
            raise RuntimeError(
                f"Type annotation should not contain calls, but '{stmt}' does"
            )
    return eval(bytecode, glob, loc)  # type: ignore[arg-type] # noqa: P204


def parse_type_line(type_line, rcb, loc):
    """Parses a type annotation specified as a comment.
    Example inputs:
        # type: (Tensor, torch.Tensor) -> Tuple[Tensor]
        # type: (Tensor, Tuple[Tensor, Tensor]) -> Tensor
    """
    arg_ann_str, ret_ann_str = split_type_line(type_line)

    try:
        arg_ann = _eval_no_call(arg_ann_str, {}, EvalEnv(rcb))
    except (NameError, SyntaxError) as e:
        raise RuntimeError(
            "Failed to parse the argument list of a type annotation"
        ) from e

    if not isinstance(arg_ann, tuple):
        arg_ann = (arg_ann,)

    try:
        ret_ann = _eval_no_call(ret_ann_str, {}, EvalEnv(rcb))
    except (NameError, SyntaxError) as e:
        raise RuntimeError(
            "Failed to parse the return type of a type annotation"
        ) from e

    arg_types = [ann_to_type(ann, loc) for ann in arg_ann]
    return arg_types, ann_to_type(ret_ann, loc)


def get_type_line(source):
    """Tries to find the line containing a comment with the type annotation."""
    type_comment = "# type:"

    lines = source.split("\n")
    lines = [(line_num, line) for line_num, line in enumerate(lines)]
    type_lines = list(filter(lambda line: type_comment in line[1], lines))
    # `type: ignore` comments may be needed in JIT'ed functions for mypy, due
    # to the hack in torch/_VF.py.

    # An ignore type comment can be of following format:
    #   1) type: ignore
    #   2) type: ignore[rule-code]
    # This ignore statement must be at the end of the line

    # adding an extra backslash before the space, to avoid triggering
    # one of the checks in .github/workflows/lint.yml
    type_pattern = re.compile("# type:\\ ignore(\\[[a-zA-Z-]+\\])?$")
    type_lines = list(filter(lambda line: not type_pattern.search(line[1]), type_lines))

    if len(type_lines) == 0:
        # Catch common typo patterns like extra spaces, typo in 'ignore', etc.
        wrong_type_pattern = re.compile("#[\t ]*type[\t ]*(?!: ignore(\\[.*\\])?$):")
        wrong_type_lines = list(
            filter(lambda line: wrong_type_pattern.search(line[1]), lines)
        )
        if len(wrong_type_lines) > 0:
            raise RuntimeError(
                "The annotation prefix in line "
                + str(wrong_type_lines[0][0])
                + " is probably invalid.\nIt must be '# type:'"
                + "\nSee PEP 484 (https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)"  # noqa: B950
                + "\nfor examples"
            )
        return None
    elif len(type_lines) == 1:
        # Only 1 type line, quit now
        return type_lines[0][1].strip()

    # Parse split up argument types according to PEP 484
    # https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code
    return_line = None
    parameter_type_lines = []
    for line_num, line in type_lines:
        if "# type: (...) -> " in line:
            return_line = (line_num, line)
            break
        elif type_comment in line:
            parameter_type_lines.append(line)
    if return_line is None:
        raise RuntimeError(
            "Return type line '# type: (...) -> ...' not found on multiline "
            "type annotation\nfor type lines:\n"
            + "\n".join([line[1] for line in type_lines])
            + "\n(See PEP 484 https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)"
        )

    def get_parameter_type(line):
        item_type = line[line.find(type_comment) + len(type_comment) :]
        return item_type.strip()

    types = map(get_parameter_type, parameter_type_lines)
    parameter_types = ", ".join(types)

    return return_line[1].replace("...", parameter_types)


def split_type_line(type_line):
    """Splits the comment with the type annotation into parts for argument and return types.
    For example, for an input of:
        # type: (Tensor, torch.Tensor) -> Tuple[Tensor, Tensor]
    This function will return:
        ("(Tensor, torch.Tensor)", "Tuple[Tensor, Tensor]")
    """
    start_offset = len("# type:")
    try:
        arrow_pos = type_line.index("->")
    except ValueError:
        raise RuntimeError(
            "Syntax error in type annotation (cound't find `->`)"
        ) from None
    return type_line[start_offset:arrow_pos].strip(), type_line[arrow_pos + 2 :].strip()


def try_real_annotations(fn, loc):
    """Tries to use the Py3.5+ annotation syntax to get the type."""
    try:
        # Note: anything annotated as `Optional[T]` will automatically
        # be returned as `Union[T, None]` per
        # https://github.com/python/typing/blob/master/src/typing.py#L850
        sig = inspect.signature(fn)
    except ValueError:
        return None

    all_annots = [sig.return_annotation] + [
        p.annotation for p in sig.parameters.values()
    ]
    if all(ann is sig.empty for ann in all_annots):
        return None

    arg_types = [ann_to_type(p.annotation, loc) for p in sig.parameters.values()]
    return_type = ann_to_type(sig.return_annotation, loc)
    return arg_types, return_type


# Finds common type for enum values belonging to an Enum class. If not all
# values have the same type, AnyType is returned.
def get_enum_value_type(e: Type[enum.Enum], loc):
    enum_values: List[enum.Enum] = list(e)
    if not enum_values:
        raise ValueError(f"No enum values defined for: '{e.__class__}'")

    types = {type(v.value) for v in enum_values}
    ir_types = [try_ann_to_type(t, loc) for t in types]

    # If Enum values are of different types, an exception will be raised here.
    # Even though Python supports this case, we chose to not implement it to
    # avoid overcomplicate logic here for a rare use case. Please report a
    # feature request if you find it necessary.
    res = torch._C.unify_type_list(ir_types)
    if not res:
        return AnyType.get()
    return res


def is_tensor(ann):
    if issubclass(ann, torch.Tensor):
        return True

    if issubclass(
        ann,
        (
            torch.LongTensor,
            torch.DoubleTensor,
            torch.FloatTensor,
            torch.IntTensor,
            torch.ShortTensor,
            torch.HalfTensor,
            torch.CharTensor,
            torch.ByteTensor,
            torch.BoolTensor,
        ),
    ):
        warnings.warn(
            "TorchScript will treat type annotations of Tensor "
            "dtype-specific subtypes as if they are normal Tensors. "
            "dtype constraints are not enforced in compilation either."
        )
        return True

    return False


def _fake_rcb(inp):
    return None


def try_ann_to_type(ann, loc, rcb=None):
    if ann is inspect.Signature.empty:
        return TensorType.getInferred()
    if ann is None:
        return NoneType.get()
    if inspect.isclass(ann) and is_tensor(ann):
        return TensorType.get()
    if is_tuple(ann):
        # Special case for the empty Tuple type annotation `Tuple[()]`
        if len(ann.__args__) == 1 and ann.__args__[0] == ():
            return TupleType([])
        return TupleType([try_ann_to_type(a, loc) for a in ann.__args__])
    if is_list(ann):
        elem_type = try_ann_to_type(ann.__args__[0], loc)
        if elem_type:
            return ListType(elem_type)
    if is_dict(ann):
        key = try_ann_to_type(ann.__args__[0], loc)
        value = try_ann_to_type(ann.__args__[1], loc)
        # Raise error if key or value is None
        if key is None:
            raise ValueError(
                f"Unknown type annotation: '{ann.__args__[0]}' at {loc.highlight()}"
            )
        if value is None:
            raise ValueError(
                f"Unknown type annotation: '{ann.__args__[1]}' at {loc.highlight()}"
            )
        return DictType(key, value)
    if is_optional(ann):
        if issubclass(ann.__args__[1], type(None)):
            contained = ann.__args__[0]
        else:
            contained = ann.__args__[1]
        valid_type = try_ann_to_type(contained, loc)
        msg = "Unsupported annotation {} could not be resolved because {} could not be resolved. At\n{}"
        assert valid_type, msg.format(repr(ann), repr(contained), repr(loc))
        return OptionalType(valid_type)
    if is_union(ann):
        # TODO: this is hack to recognize NumberType
        if set(ann.__args__) == {int, float, complex}:
            return NumberType.get()
        inner: List = []
        # We need these extra checks because both `None` and invalid
        # values will return `None`
        # TODO: Determine if the other cases need to be fixed as well
        for a in ann.__args__:
            if a is None:
                inner.append(NoneType.get())
            maybe_type = try_ann_to_type(a, loc)
            msg = "Unsupported annotation {} could not be resolved because {} could not be resolved. At\n{}"
            assert maybe_type, msg.format(repr(ann), repr(maybe_type), repr(loc))
            inner.append(maybe_type)
        return UnionType(inner)  # type: ignore[arg-type]
    if torch.distributed.rpc.is_available() and is_rref(ann):
        return RRefType(try_ann_to_type(ann.__args__[0], loc))
    if is_future(ann):
        return FutureType(try_ann_to_type(ann.__args__[0], loc))
    if is_await(ann):
        elementType = (
            try_ann_to_type(ann.__args__[0], loc)
            if hasattr(ann, "__args__")
            else AnyType.get()
        )
        return AwaitType(elementType)
    if ann is float:
        return FloatType.get()
    if ann is complex:
        return ComplexType.get()
    if ann is int:
        return IntType.get()
    if ann is str:
        return StringType.get()
    if ann is bool:
        return BoolType.get()
    if ann is Any:
        return AnyType.get()
    if ann is type(None):
        return NoneType.get()
    if inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"):
        return InterfaceType(ann.__torch_script_interface__)
    if ann is torch.device:
        return DeviceObjType.get()
    if ann is torch.Stream:
        return StreamObjType.get()
    if ann is torch.dtype:
        return IntType.get()  # dtype not yet bound in as its own type
    if inspect.isclass(ann) and issubclass(ann, enum.Enum):
        if _get_script_class(ann) is None:
            scripted_class = torch.jit._script._recursive_compile_class(ann, loc)
            name = scripted_class.qualified_name()
        else:
            name = _qualified_name(ann)
        return EnumType(name, get_enum_value_type(ann, loc), list(ann))
    if inspect.isclass(ann):
        maybe_script_class = _get_script_class(ann)
        if maybe_script_class is not None:
            return maybe_script_class
        if torch._jit_internal.can_compile_class(ann):
            return torch.jit._script._recursive_compile_class(ann, loc)

    # Maybe resolve a NamedTuple to a Tuple Type
    if rcb is None:
        rcb = _fake_rcb
    return torch._C._resolve_type_from_object(ann, loc, rcb)


def ann_to_type(ann, loc, rcb=None):
    the_type = try_ann_to_type(ann, loc, rcb)
    if the_type is not None:
        return the_type
    raise ValueError(f"Unknown type annotation: '{ann}' at {loc.highlight()}")


__all__ = [
    "Any",
    "List",
    "BroadcastingList1",
    "BroadcastingList2",
    "BroadcastingList3",
    "Tuple",
    "is_tuple",
    "is_list",
    "Dict",
    "is_dict",
    "is_optional",
    "is_union",
    "TensorType",
    "TupleType",
    "FloatType",
    "ComplexType",
    "IntType",
    "ListType",
    "StringType",
    "DictType",
    "AnyType",
    "Module",
    # TODO: Consider not exporting these during wildcard import (reserve
    # that for the types; for idiomatic typing code.)
    "get_signature",
    "check_fn",
    "get_param_names",
    "parse_type_line",
    "get_type_line",
    "split_type_line",
    "try_real_annotations",
    "try_ann_to_type",
    "ann_to_type",
]

	import ast
	import builtins
	import dis
	import enum
	import inspect
	import re
	import warnings

	from textwrap import dedent
	from typing import Type

	import torch

	from torch._C import (
	AnyType,
	AwaitType,
	BoolType,
	ComplexType,
	DeviceObjType,
	DictType,
	EnumType,
	FloatType,
	FutureType,
	InterfaceType,
	IntType,
	ListType,
	NoneType,
	NumberType,
	OptionalType,
	StreamObjType,
	StringType,
	TensorType,
	TupleType,
	UnionType,
	)
	from torch._sources import get_source_lines_and_file
	from .._jit_internal import ( # type: ignore[attr-defined]
	_Await,
	_qualified_name,
	Any,
	BroadcastingList1,
	BroadcastingList2,
	BroadcastingList3,
	Dict,
	Future,
	is_await,
	is_dict,
	is_future,
	is_ignored_fn,
	is_list,
	is_optional,
	is_tuple,
	is_union,
	List,
	Optional,
	Tuple,
	Union,
	)
	from ._state import _get_script_class

	if torch.distributed.rpc.is_available():
	from torch._C import RRefType
	from .._jit_internal import is_rref, RRef

	from torch._ops import OpOverloadPacket


	class Module:
	def __init__(self, name, members):
	self.name = name
	self.members = members

	def __getattr__(self, name):
	try:
	return self.members[name]
	except KeyError:
	raise RuntimeError(
	f"Module {self.name} has no member called {name}"
	) from None


	class EvalEnv:
	env = {
	"torch": Module("torch", {"Tensor": torch.Tensor}),
	"Tensor": torch.Tensor,
	"typing": Module("typing", {"Tuple": Tuple}),
	"Tuple": Tuple,
	"List": List,
	"Dict": Dict,
	"Optional": Optional,
	"Union": Union,
	"Future": Future,
	"Await": _Await,
	}

	def __init__(self, rcb):
	self.rcb = rcb
	if torch.distributed.rpc.is_available():
	self.env["RRef"] = RRef

	def __getitem__(self, name):
	if name in self.env:
	return self.env[name]
	if self.rcb is not None:
	return self.rcb(name)
	return getattr(builtins, name, None)


	def get_signature(fn, rcb, loc, is_method):
	if isinstance(fn, OpOverloadPacket):
	signature = try_real_annotations(fn.op, loc)
	else:
	signature = try_real_annotations(fn, loc)
	if signature is not None and is_method:
	# If this is a method, then the signature will include a type for
	# `self`, but type comments do not contain a `self`. So strip it
	# away here so everything is consistent (`inspect.ismethod` does
	# not work here since `fn` is unbound at this point)
	param_types, return_type = signature
	param_types = param_types[1:]
	signature = (param_types, return_type)

	if signature is None:
	type_line, source = None, None
	try:
	source = dedent("".join(get_source_lines_and_file(fn)[0]))
	type_line = get_type_line(source)
	except TypeError:
	pass
	# This might happen both because we failed to get the source of fn, or
	# because it didn't have any annotations.
	if type_line is not None:
	signature = parse_type_line(type_line, rcb, loc)

	return signature


	def is_function_or_method(the_callable):
	# A stricter version of `inspect.isroutine` that does not pass for built-in
	# functions
	return inspect.isfunction(the_callable) or inspect.ismethod(the_callable)


	def is_vararg(the_callable):
	if not is_function_or_method(the_callable) and callable(the_callable): # noqa: B004
	# If `the_callable` is a class, de-sugar the call so we can still get
	# the signature
	the_callable = the_callable.__call__

	if is_function_or_method(the_callable):
	return inspect.getfullargspec(the_callable).varargs is not None
	else:
	return False


	def get_param_names(fn, n_args):
	if isinstance(fn, OpOverloadPacket):
	fn = fn.op

	if (
	not is_function_or_method(fn)
	and callable(fn)
	and is_function_or_method(fn.__call__)
	): # noqa: B004
	# De-sugar calls to classes
	fn = fn.__call__

	if is_function_or_method(fn):
	if is_ignored_fn(fn):
	fn = inspect.unwrap(fn)
	return inspect.getfullargspec(fn).args
	else:
	# The `fn` was not a method or function (maybe a class with a __call__
	# method, so use a default param name list)
	return [str(i) for i in range(n_args)]


	def check_fn(fn, loc):
	# Make sure the function definition is not a class instantiation
	try:
	source = dedent("".join(get_source_lines_and_file(fn)[0]))
	except (OSError, TypeError):
	return
	if source is None:
	return

	py_ast = ast.parse(source)
	if len(py_ast.body) == 1 and isinstance(py_ast.body[0], ast.ClassDef):
	raise torch.jit.frontend.FrontendError(
	loc,
	f"Cannot instantiate class '{py_ast.body[0].name}' in a script function",
	)
	if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef):
	raise torch.jit.frontend.FrontendError(
	loc, "Expected a single top-level function"
	)


	def _eval_no_call(stmt, glob, loc):
	"""Evaluate statement as long as it does not contain any method/function calls"""
	bytecode = compile(stmt, "", mode="eval")
	for insn in dis.get_instructions(bytecode):
	if "CALL" in insn.opname:
	raise RuntimeError(
	f"Type annotation should not contain calls, but '{stmt}' does"
	)
	return eval(bytecode, glob, loc) # type: ignore[arg-type] # noqa: P204


	def parse_type_line(type_line, rcb, loc):
	"""Parses a type annotation specified as a comment.

	Example inputs:
	# type: (Tensor, torch.Tensor) -> Tuple[Tensor]
	# type: (Tensor, Tuple[Tensor, Tensor]) -> Tensor
	"""
	arg_ann_str, ret_ann_str = split_type_line(type_line)

	try:
	arg_ann = _eval_no_call(arg_ann_str, {}, EvalEnv(rcb))
	except (NameError, SyntaxError) as e:
	raise RuntimeError(
	"Failed to parse the argument list of a type annotation"
	) from e

	if not isinstance(arg_ann, tuple):
	arg_ann = (arg_ann,)

	try:
	ret_ann = _eval_no_call(ret_ann_str, {}, EvalEnv(rcb))
	except (NameError, SyntaxError) as e:
	raise RuntimeError(
	"Failed to parse the return type of a type annotation"
	) from e

	arg_types = [ann_to_type(ann, loc) for ann in arg_ann]
	return arg_types, ann_to_type(ret_ann, loc)


	def get_type_line(source):
	"""Tries to find the line containing a comment with the type annotation."""
	type_comment = "# type:"

	lines = source.split("\n")
	lines = [(line_num, line) for line_num, line in enumerate(lines)]
	type_lines = list(filter(lambda line: type_comment in line[1], lines))
	# `type: ignore` comments may be needed in JIT'ed functions for mypy, due
	# to the hack in torch/_VF.py.

	# An ignore type comment can be of following format:
	# 1) type: ignore
	# 2) type: ignore[rule-code]
	# This ignore statement must be at the end of the line

	# adding an extra backslash before the space, to avoid triggering
	# one of the checks in .github/workflows/lint.yml
	type_pattern = re.compile("# type:\\ ignore(\\[[a-zA-Z-]+\\])?$")
	type_lines = list(filter(lambda line: not type_pattern.search(line[1]), type_lines))

	if len(type_lines) == 0:
	# Catch common typo patterns like extra spaces, typo in 'ignore', etc.
	wrong_type_pattern = re.compile("#[\t ]type[\t ](?!: ignore(\\[.*\\])?$):")
	wrong_type_lines = list(
	filter(lambda line: wrong_type_pattern.search(line[1]), lines)
	)
	if len(wrong_type_lines) > 0:
	raise RuntimeError(
	"The annotation prefix in line "
	+ str(wrong_type_lines[0][0])
	+ " is probably invalid.\nIt must be '# type:'"
	+ "\nSee PEP 484 (https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)" # noqa: B950
	+ "\nfor examples"
	)
	return None
	elif len(type_lines) == 1:
	# Only 1 type line, quit now
	return type_lines[0][1].strip()

	# Parse split up argument types according to PEP 484
	# https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code
	return_line = None
	parameter_type_lines = []
	for line_num, line in type_lines:
	if "# type: (...) -> " in line:
	return_line = (line_num, line)
	break
	elif type_comment in line:
	parameter_type_lines.append(line)
	if return_line is None:
	raise RuntimeError(
	"Return type line '# type: (...) -> ...' not found on multiline "
	"type annotation\nfor type lines:\n"
	+ "\n".join([line[1] for line in type_lines])
	+ "\n(See PEP 484 https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)"
	)

	def get_parameter_type(line):
	item_type = line[line.find(type_comment) + len(type_comment) :]
	return item_type.strip()

	types = map(get_parameter_type, parameter_type_lines)
	parameter_types = ", ".join(types)

	return return_line[1].replace("...", parameter_types)


	def split_type_line(type_line):
	"""Splits the comment with the type annotation into parts for argument and return types.

	For example, for an input of:
	# type: (Tensor, torch.Tensor) -> Tuple[Tensor, Tensor]

	This function will return:
	("(Tensor, torch.Tensor)", "Tuple[Tensor, Tensor]")

	"""
	start_offset = len("# type:")
	try:
	arrow_pos = type_line.index("->")
	except ValueError:
	raise RuntimeError(
	"Syntax error in type annotation (cound't find `->`)"
	) from None
	return type_line[start_offset:arrow_pos].strip(), type_line[arrow_pos + 2 :].strip()


	def try_real_annotations(fn, loc):
	"""Tries to use the Py3.5+ annotation syntax to get the type."""
	try:
	# Note: anything annotated as `Optional[T]` will automatically
	# be returned as `Union[T, None]` per
	# https://github.com/python/typing/blob/master/src/typing.py#L850
	sig = inspect.signature(fn)
	except ValueError:
	return None

	all_annots = [sig.return_annotation] + [
	p.annotation for p in sig.parameters.values()
	]
	if all(ann is sig.empty for ann in all_annots):
	return None

	arg_types = [ann_to_type(p.annotation, loc) for p in sig.parameters.values()]
	return_type = ann_to_type(sig.return_annotation, loc)
	return arg_types, return_type


	# Finds common type for enum values belonging to an Enum class. If not all
	# values have the same type, AnyType is returned.
	def get_enum_value_type(e: Type[enum.Enum], loc):
	enum_values: List[enum.Enum] = list(e)
	if not enum_values:
	raise ValueError(f"No enum values defined for: '{e.__class__}'")

	types = {type(v.value) for v in enum_values}
	ir_types = [try_ann_to_type(t, loc) for t in types]

	# If Enum values are of different types, an exception will be raised here.
	# Even though Python supports this case, we chose to not implement it to
	# avoid overcomplicate logic here for a rare use case. Please report a
	# feature request if you find it necessary.
	res = torch._C.unify_type_list(ir_types)
	if not res:
	return AnyType.get()
	return res


	def is_tensor(ann):
	if issubclass(ann, torch.Tensor):
	return True

	if issubclass(
	ann,
	(
	torch.LongTensor,
	torch.DoubleTensor,
	torch.FloatTensor,
	torch.IntTensor,
	torch.ShortTensor,
	torch.HalfTensor,
	torch.CharTensor,
	torch.ByteTensor,
	torch.BoolTensor,
	),
	):
	warnings.warn(
	"TorchScript will treat type annotations of Tensor "
	"dtype-specific subtypes as if they are normal Tensors. "
	"dtype constraints are not enforced in compilation either."
	)
	return True

	return False


	def _fake_rcb(inp):
	return None


	def try_ann_to_type(ann, loc, rcb=None):
	if ann is inspect.Signature.empty:
	return TensorType.getInferred()
	if ann is None:
	return NoneType.get()
	if inspect.isclass(ann) and is_tensor(ann):
	return TensorType.get()
	if is_tuple(ann):
	# Special case for the empty Tuple type annotation `Tuple[()]`
	if len(ann.__args__) == 1 and ann.__args__[0] == ():
	return TupleType([])
	return TupleType([try_ann_to_type(a, loc) for a in ann.__args__])
	if is_list(ann):
	elem_type = try_ann_to_type(ann.__args__[0], loc)
	if elem_type:
	return ListType(elem_type)
	if is_dict(ann):
	key = try_ann_to_type(ann.__args__[0], loc)
	value = try_ann_to_type(ann.__args__[1], loc)
	# Raise error if key or value is None
	if key is None:
	raise ValueError(
	f"Unknown type annotation: '{ann.__args__[0]}' at {loc.highlight()}"
	)
	if value is None:
	raise ValueError(
	f"Unknown type annotation: '{ann.__args__[1]}' at {loc.highlight()}"
	)
	return DictType(key, value)
	if is_optional(ann):
	if issubclass(ann.__args__[1], type(None)):
	contained = ann.__args__[0]
	else:
	contained = ann.__args__[1]
	valid_type = try_ann_to_type(contained, loc)
	msg = "Unsupported annotation {} could not be resolved because {} could not be resolved. At\n{}"
	assert valid_type, msg.format(repr(ann), repr(contained), repr(loc))
	return OptionalType(valid_type)
	if is_union(ann):
	# TODO: this is hack to recognize NumberType
	if set(ann.__args__) == {int, float, complex}:
	return NumberType.get()
	inner: List = []
	# We need these extra checks because both `None` and invalid
	# values will return `None`
	# TODO: Determine if the other cases need to be fixed as well
	for a in ann.__args__:
	if a is None:
	inner.append(NoneType.get())
	maybe_type = try_ann_to_type(a, loc)
	msg = "Unsupported annotation {} could not be resolved because {} could not be resolved. At\n{}"
	assert maybe_type, msg.format(repr(ann), repr(maybe_type), repr(loc))
	inner.append(maybe_type)
	return UnionType(inner) # type: ignore[arg-type]
	if torch.distributed.rpc.is_available() and is_rref(ann):
	return RRefType(try_ann_to_type(ann.__args__[0], loc))
	if is_future(ann):
	return FutureType(try_ann_to_type(ann.__args__[0], loc))
	if is_await(ann):
	elementType = (
	try_ann_to_type(ann.__args__[0], loc)
	if hasattr(ann, "__args__")
	else AnyType.get()
	)
	return AwaitType(elementType)
	if ann is float:
	return FloatType.get()
	if ann is complex:
	return ComplexType.get()
	if ann is int:
	return IntType.get()
	if ann is str:
	return StringType.get()
	if ann is bool:
	return BoolType.get()
	if ann is Any:
	return AnyType.get()
	if ann is type(None):
	return NoneType.get()
	if inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"):
	return InterfaceType(ann.__torch_script_interface__)
	if ann is torch.device:
	return DeviceObjType.get()
	if ann is torch.Stream:
	return StreamObjType.get()
	if ann is torch.dtype:
	return IntType.get() # dtype not yet bound in as its own type
	if inspect.isclass(ann) and issubclass(ann, enum.Enum):
	if _get_script_class(ann) is None:
	scripted_class = torch.jit._script._recursive_compile_class(ann, loc)
	name = scripted_class.qualified_name()
	else:
	name = _qualified_name(ann)
	return EnumType(name, get_enum_value_type(ann, loc), list(ann))
	if inspect.isclass(ann):
	maybe_script_class = _get_script_class(ann)
	if maybe_script_class is not None:
	return maybe_script_class
	if torch._jit_internal.can_compile_class(ann):
	return torch.jit._script._recursive_compile_class(ann, loc)

	# Maybe resolve a NamedTuple to a Tuple Type
	if rcb is None:
	rcb = _fake_rcb
	return torch._C._resolve_type_from_object(ann, loc, rcb)


	def ann_to_type(ann, loc, rcb=None):
	the_type = try_ann_to_type(ann, loc, rcb)
	if the_type is not None:
	return the_type
	raise ValueError(f"Unknown type annotation: '{ann}' at {loc.highlight()}")


	__all__ = [
	"Any",
	"List",
	"BroadcastingList1",
	"BroadcastingList2",
	"BroadcastingList3",
	"Tuple",
	"is_tuple",
	"is_list",
	"Dict",
	"is_dict",
	"is_optional",
	"is_union",
	"TensorType",
	"TupleType",
	"FloatType",
	"ComplexType",
	"IntType",
	"ListType",
	"StringType",
	"DictType",
	"AnyType",
	"Module",
	# TODO: Consider not exporting these during wildcard import (reserve
	# that for the types; for idiomatic typing code.)
	"get_signature",
	"check_fn",
	"get_param_names",
	"parse_type_line",
	"get_type_line",
	"split_type_line",
	"try_real_annotations",
	"try_ann_to_type",
	"ann_to_type",
	]