Spaces:
Sleeping
Sleeping
| import ast | |
| import inspect | |
| import textwrap | |
| import warnings | |
| import torch | |
| class AttributeTypeIsSupportedChecker(ast.NodeVisitor): | |
| """Check the ``__init__`` method of a given ``nn.Module``. | |
| It ensures that all instance-level attributes can be properly initialized. | |
| Specifically, we do type inference based on attribute values...even | |
| if the attribute in question has already been typed using | |
| Python3-style annotations or ``torch.jit.annotate``. This means that | |
| setting an instance-level attribute to ``[]`` (for ``List``), | |
| ``{}`` for ``Dict``), or ``None`` (for ``Optional``) isn't enough | |
| information for us to properly initialize that attribute. | |
| An object of this class can walk a given ``nn.Module``'s AST and | |
| determine if it meets our requirements or not. | |
| Known limitations | |
| 1. We can only check the AST nodes for certain constructs; we can't | |
| ``eval`` arbitrary expressions. This means that function calls, | |
| class instantiations, and complex expressions that resolve to one of | |
| the "empty" values specified above will NOT be flagged as | |
| problematic. | |
| 2. We match on string literals, so if the user decides to use a | |
| non-standard import (e.g. `from typing import List as foo`), we | |
| won't catch it. | |
| Example: | |
| .. code-block:: python | |
| class M(torch.nn.Module): | |
| def fn(self): | |
| return [] | |
| def __init__(self): | |
| super().__init__() | |
| self.x: List[int] = [] | |
| def forward(self, x: List[int]): | |
| self.x = x | |
| return 1 | |
| The above code will pass the ``AttributeTypeIsSupportedChecker`` | |
| check since we have a function call in ``__init__``. However, | |
| it will still fail later with the ``RuntimeError`` "Tried to set | |
| nonexistent attribute: x. Did you forget to initialize it in | |
| __init__()?". | |
| Args: | |
| nn_module - The instance of ``torch.nn.Module`` whose | |
| ``__init__`` method we wish to check | |
| """ | |
| def check(self, nn_module: torch.nn.Module) -> None: | |
| source_lines = inspect.getsource(nn_module.__class__.__init__) | |
| # Ignore comments no matter the indentation | |
| def is_useless_comment(line): | |
| line = line.strip() | |
| return line.startswith("#") and not line.startswith("# type:") | |
| source_lines = "\n".join( | |
| [l for l in source_lines.split("\n") if not is_useless_comment(l)] | |
| ) | |
| # This AST only contains the `__init__` method of the nn.Module | |
| init_ast = ast.parse(textwrap.dedent(source_lines)) | |
| # Get items annotated in the class body | |
| self.class_level_annotations = list(nn_module.__annotations__.keys()) | |
| # Flag for later | |
| self.visiting_class_level_ann = False | |
| self.visit(init_ast) | |
| def _is_empty_container(self, node: ast.AST, ann_type: str) -> bool: | |
| if ann_type == "List": | |
| # Assigning `[]` to a `List` type gives you a Node where | |
| # value=List(elts=[], ctx=Load()) | |
| if not isinstance(node, ast.List): | |
| return False | |
| if node.elts: | |
| return False | |
| elif ann_type == "Dict": | |
| # Assigning `{}` to a `Dict` type gives you a Node where | |
| # value=Dict(keys=[], values=[]) | |
| if not isinstance(node, ast.Dict): | |
| return False | |
| if node.keys: | |
| return False | |
| elif ann_type == "Optional": | |
| # Assigning `None` to an `Optional` type gives you a | |
| # Node where value=Constant(value=None, kind=None) | |
| if not isinstance(node, ast.Constant): | |
| return False | |
| if node.value: # type: ignore[attr-defined] | |
| return False | |
| return True | |
| def visit_Assign(self, node): | |
| """Store assignment state when assigning to a Call Node. | |
| If we're visiting a Call Node (the right-hand side of an | |
| assignment statement), we won't be able to check the variable | |
| that we're assigning to (the left-hand side of an assignment). | |
| Because of this, we need to store this state in visitAssign. | |
| (Luckily, we only have to do this if we're assigning to a Call | |
| Node, i.e. ``torch.jit.annotate``. If we're using normal Python | |
| annotations, we'll be visiting an AnnAssign Node, which has its | |
| target built in.) | |
| """ | |
| try: | |
| if ( | |
| isinstance(node.value, ast.Call) | |
| and node.targets[0].attr in self.class_level_annotations | |
| ): | |
| self.visiting_class_level_ann = True | |
| except AttributeError: | |
| return | |
| self.generic_visit(node) | |
| self.visiting_class_level_ann = False | |
| def visit_AnnAssign(self, node): | |
| """Visit an AnnAssign node in an ``nn.Module``'s ``__init__`` method. | |
| It checks if it conforms to our attribute annotation rules.""" | |
| # If we have a local variable | |
| try: | |
| if node.target.value.id != "self": | |
| return | |
| except AttributeError: | |
| return | |
| # If we have an attribute that's already been annotated at the | |
| # class level | |
| if node.target.attr in self.class_level_annotations: | |
| return | |
| # TODO @ansley: add `Union` once landed | |
| # NB: Even though `Tuple` is a "container", we don't want to | |
| # check for it here. `Tuple` functions as an type with an | |
| # "infinite" number of subtypes, in the sense that you can have | |
| # `Tuple[())]`, `Tuple[T1]`, `Tuple[T2]`, `Tuple[T1, T2]`, | |
| # `Tuple[T2, T1]` and so on, and none of these subtypes can be | |
| # used in place of the other. Therefore, assigning an empty | |
| # tuple in `__init__` CORRECTLY means that that variable | |
| # cannot be reassigned later to a non-empty tuple. Same | |
| # deal with `NamedTuple` | |
| containers = {"List", "Dict", "Optional"} | |
| # If we're not evaluating one of the specified problem types | |
| try: | |
| if node.annotation.value.id not in containers: | |
| return | |
| except AttributeError: | |
| # To evaluate a base type (`str`, `int`, etc.), we would | |
| # have needed to get the name through `node.annotation.id` | |
| # instead of `node.annotation.value.id`. Seems that we're | |
| # not evaluating one of our "containers" | |
| return | |
| # Check if the assigned variable is empty | |
| ann_type = node.annotation.value.id | |
| if not self._is_empty_container(node.value, ann_type): | |
| return | |
| warnings.warn( | |
| "The TorchScript type system doesn't support " | |
| "instance-level annotations on empty non-base " | |
| "types in `__init__`. Instead, either 1) use a " | |
| "type annotation in the class body, or 2) wrap " | |
| "the type in `torch.jit.Attribute`." | |
| ) | |
| def visit_Call(self, node): | |
| """Determine if a Call node is 'torch.jit.annotate' in __init__. | |
| Visit a Call node in an ``nn.Module``'s ``__init__`` | |
| method and determine if it's ``torch.jit.annotate``. If so, | |
| see if it conforms to our attribute annotation rules. | |
| """ | |
| # If we have an attribute that's already been annotated at the | |
| # class level | |
| if self.visiting_class_level_ann: | |
| return | |
| # If this isn't a call to `torch.jit.annotate` | |
| try: | |
| if ( | |
| node.func.value.value.id != "torch" | |
| or node.func.value.attr != "jit" | |
| or node.func.attr != "annotate" | |
| ): | |
| self.generic_visit(node) | |
| elif ( | |
| node.func.value.value.id != "jit" or node.func.value.attr != "annotate" | |
| ): | |
| self.generic_visit(node) | |
| except AttributeError: | |
| # Looks like we didn't even have the right node structure | |
| # to check for `torch.jit.annotate` in the first place | |
| self.generic_visit(node) | |
| # Invariant: we have a `torch.jit.annotate` or a | |
| # `torch.annotate` call | |
| # A Call Node for `torch.jit.annotate` should have an `args` | |
| # list of length 2 where args[0] represents the annotation and | |
| # args[1] represents the actual value | |
| if len(node.args) != 2: | |
| return | |
| if not isinstance(node.args[0], ast.Subscript): | |
| return | |
| # See notes in `visit_AnnAssign` r.e. containers | |
| containers = {"List", "Dict", "Optional"} | |
| try: | |
| ann_type = node.args[0].value.id # type: ignore[attr-defined] | |
| except AttributeError: | |
| return | |
| if ann_type not in containers: | |
| return | |
| # Check if the assigned variable is empty | |
| if not self._is_empty_container(node.args[1], ann_type): | |
| return | |
| warnings.warn( | |
| "The TorchScript type system doesn't support " | |
| "instance-level annotations on empty non-base " | |
| "types in `__init__`. Instead, either 1) use a " | |
| "type annotation in the class body, or 2) wrap " | |
| "the type in `torch.jit.Attribute`." | |
| ) | |