[inductor] Enable mypy checking in lowering.py (pytorch#105317)

Summary:

As suggested in pytorch#105230, mypy checking is enabled in `torch/_inductor/lowering.py`.

23 errors fixed; 6 silenced with `# type: ignore[attr-defined]`.

Test Plan:

Before the fix:

```
$ mypy torch/_inductor/lowering.py

torch/_inductor/lowering.py:139:16: error: "Symbol" has no attribute "is_integer"  [attr-defined]
torch/_inductor/lowering.py:263:20: error: Incompatible types in assignment (expression has type "Union[List[Any], Tuple[Any, ...]]", variable has type "List[Any]")  [assignment]
torch/_inductor/lowering.py:427:49: error: "IRNode" has no attribute "get_size"  [attr-defined]
torch/_inductor/lowering.py:439:37: error: "IRNode" has no attribute "get_dtype"  [attr-defined]
torch/_inductor/lowering.py:456:34: error: "IRNode" has no attribute "get_device"  [attr-defined]
torch/_inductor/lowering.py:645:44: error: Need type annotation for "b"  [var-annotated]
torch/_inductor/lowering.py:1321:12: error: "FakeTensor" has no attribute "is_cpu"  [attr-defined]
torch/_inductor/lowering.py:1542:24: error: Argument 3 to "FixedLayout" has incompatible type "List[int]"; expected "List[Expr]"  [arg-type]
torch/_inductor/lowering.py:1542:81: error: Argument "offset" to "FixedLayout" has incompatible type "int"; expected "Expr"  [arg-type]
torch/_inductor/lowering.py:1571:24: error: Argument 3 to "FixedLayout" has incompatible type "List[int]"; expected "List[Expr]"  [arg-type]
torch/_inductor/lowering.py:1571:81: error: Argument "offset" to "FixedLayout" has incompatible type "int"; expected "Expr"  [arg-type]
torch/_inductor/lowering.py:1654:12: error: Incompatible types in assignment (expression has type "List[Any]", variable has type "Tuple[Any, ...]")  [assignment]
torch/_inductor/lowering.py:2009:9: error: Need type annotation for "ranges" (hint: "ranges: List[<type>] = ...")  [var-annotated]
torch/_inductor/lowering.py:2151:16: error: Incompatible types in assignment (expression has type "List[Any]", variable has type "Tuple[Any, ...]")  [assignment]
torch/_inductor/lowering.py:2198:43: error: Item "type" of "Union[List[Any], type]" has no attribute "__iter__" (not iterable)  [union-attr]
torch/_inductor/lowering.py:2229:36: error: Argument 1 to "len" has incompatible type "Union[List[Any], type]"; expected "Sized"  [arg-type]
torch/_inductor/lowering.py:2231:38: error: Item "type" of "Union[List[Any], type]" has no attribute "__iter__" (not iterable)  [union-attr]
torch/_inductor/lowering.py:2233:35: error: Item "type" of "Union[List[Any], type]" has no attribute "__iter__" (not iterable)  [union-attr]
torch/_inductor/lowering.py:2569:54: error: Incompatible default for argument "reduce" (default has type "None", argument has type "str")  [assignment]
torch/_inductor/lowering.py:2569:54: note: PEP 484 prohibits implicit Optional. Accordingly, mypy has changed its default to no_implicit_optional=True
torch/_inductor/lowering.py:2569:54: note: Use https://github.com/hauntsaninja/no_implicit_optional to automatically upgrade your codebase
torch/_inductor/lowering.py:2586:59: error: Incompatible default for argument "reduce" (default has type "None", argument has type "str")  [assignment]
torch/_inductor/lowering.py:2586:59: note: PEP 484 prohibits implicit Optional. Accordingly, mypy has changed its default to no_implicit_optional=True
torch/_inductor/lowering.py:2586:59: note: Use https://github.com/hauntsaninja/no_implicit_optional to automatically upgrade your codebase
torch/_inductor/lowering.py:2720:65: error: Incompatible default for argument "scales_x" (default has type "None", argument has type "Tuple[float]")  [assignment]
torch/_inductor/lowering.py:2720:65: note: PEP 484 prohibits implicit Optional. Accordingly, mypy has changed its default to no_implicit_optional=True
torch/_inductor/lowering.py:2720:65: note: Use https://github.com/hauntsaninja/no_implicit_optional to automatically upgrade your codebase
torch/_inductor/lowering.py:2735:5: error: Name "scale" already defined on line 2731  [no-redef]
torch/_inductor/lowering.py:2758:47: error: Argument 3 to "upsample_nearestnd" has incompatible type "Tuple[Optional[float]]"; expected "Tuple[float]"  [arg-type]
torch/_inductor/lowering.py:2765:47: error: Argument 3 to "upsample_nearestnd" has incompatible type "Tuple[Optional[float], Optional[float]]"; expected "Tuple[float]"  [arg-type]
torch/_inductor/lowering.py:2776:47: error: Argument 3 to "upsample_nearestnd" has incompatible type "Tuple[Optional[float], Optional[float], Optional[float]]"; expected "Tuple[float]"  [arg-type]
torch/_inductor/lowering.py:2949:13: error: No binding for nonlocal "grad" found  [misc]
torch/_inductor/lowering.py:3063:49: error: Argument 2 to "range_mask_low" has incompatible type "int"; expected "Expr"  [arg-type]
torch/_inductor/lowering.py:3271:48: error: "IRNode" has no attribute "data"  [attr-defined]
torch/_inductor/lowering.py:3272:16: error: "IRNode" has no attribute "data"  [attr-defined]
Found 29 errors in 1 file (checked 1 source file)
```

After the fix:

```
$ mypy torch/_inductor/lowering.py

Success: no issues found in 1 source file
```

Reviewers: @eellison

Subscribers:

Tasks:

Tags:

Pull Request resolved: pytorch#105317
Approved by: https://github.com/eellison

Author: aakhundov

.lintrunner.toml

@@ -182,6 +182,7 @@ include_patterns = [
     'torch/_inductor/graph.py',
     'torch/_inductor/codegen/wrapper.py',
     'torch/_inductor/cudagraph_trees.py',
+    'torch/_inductor/lowering.py',
     'torch/_C/_dynamo/**/*.py',
     'test/test_utils.py',  # used to by in MYPY but after importing op_db it took 10+ minutes
 ]

tools/pyi/gen_pyi.py

@@ -1035,6 +1035,7 @@ def replace_special_case(hint: str) -> str:
                 "def is_contiguous(self, memory_format=torch.contiguous_format) -> _bool: ..."
             ],
             "_is_view": ["def _is_view(self) -> _bool: ..."],
+            "is_cpu": ["is_cpu: _bool"],
             "is_cuda": ["is_cuda: _bool"],
             "is_leaf": ["is_leaf: _bool"],
             "is_nested": ["is_nested: _bool"],

torch/_inductor/ir.py

@@ -1818,9 +1818,9 @@ def __init__(
         self,
         device: torch.device,
         dtype: torch.dtype,
-        size: List[Expr],
-        stride: List[Expr] = None,
-        offset: Expr = Integer(0),
+        size: Union[List[Expr], List[int]],
+        stride: Optional[Union[List[Expr], List[int]]] = None,
+        offset: Union[Expr, int] = Integer(0),
     ):
         if stride is None:
             stride = FlexibleLayout.contiguous_strides(size)
@@ -3116,7 +3116,7 @@ def __init__(
         index,
         src,
         *,
-        reduce: str = None,
+        reduce: Optional[str] = None,
         include_self: bool = True,
     ):
         assert fn in {"aten.scatter_", "aten.scatter_reduce_"}

torch/_inductor/lowering.py

@@ -5,7 +5,7 @@
 import warnings
 from collections import defaultdict
 from collections.abc import Iterable
-from typing import List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Union
 
 import sympy
 
@@ -130,7 +130,7 @@ def is_integer_type(x):
     if isinstance(x, TensorBox):
         return is_integer_dtype(x.get_dtype()) or is_boolean_dtype(x.get_dtype())
     elif isinstance(x, sympy.Symbol):
-        return x.is_integer is True
+        return x.is_integer is True  # type: ignore[attr-defined]
     else:
         return isinstance(x, int)
 
@@ -249,7 +249,7 @@ def _register_lowering(
 
     @functools.wraps(decomp_fn)
     def wrapped(*args, **kwargs):
-        args = list(args)
+        args: Union[List[Any], Tuple[Any, ...]] = list(args)
         unpacked = False
         # TODO maybe we need to use pytrees here
         if len(args) == 1 and isinstance(args[0], (list, tuple)):
@@ -418,7 +418,7 @@ def inner(*inputs: List[List[TensorBox]], alpha=1):
         def is_dynamic(*args):
             return any(
                 isinstance(t, TensorBox)
-                and any(x.free_symbols for x in t.data.get_size())
+                and any(x.free_symbols for x in t.data.get_size())  # type: ignore[attr-defined]
                 for t in args
             )
 
@@ -430,7 +430,7 @@ def has_type_promotion(*args):
                 for t in args:
                     if isinstance(t, TensorBox):
                         if dtype is None:
-                            dtype = t.data.get_dtype()
+                            dtype = t.data.get_dtype()  # type: ignore[attr-defined]
                         elif dtype != t.data.get_dtype():
                             return True
                 return False
@@ -447,7 +447,7 @@ def group_args(arg_pairs):
                 device = None
                 for t in args:
                     if isinstance(t, TensorBox):
-                        device = t.data.get_device()
+                        device = t.data.get_device()  # type: ignore[attr-defined]
                         break
                 assert (
                     device is not None
@@ -630,8 +630,8 @@ def fn(*args):
 def broadcast_tensors(*inputs):
     if len(inputs) == 1 and isinstance(inputs[0], (list, tuple)):
         return broadcast_tensors(*inputs[0])
-    target = functools.reduce(
-        broadcast_symbolic_shapes, [x.get_size() for x in inputs], ()
+    target: List[sympy.Expr] = functools.reduce(
+        broadcast_symbolic_shapes, [x.get_size() for x in inputs], []
     )
     outputs = []
     for x in inputs:
@@ -1645,7 +1645,9 @@ def apply_constraint(arg, fx_arg):
             return ir.ExternKernel.require_stride_order(arg, stride_order)
         return arg
 
-    args = [apply_constraint(arg, fx_arg) for arg, fx_arg in zip(args, fx_node.args)]
+    args = tuple(
+        apply_constraint(arg, fx_arg) for arg, fx_arg in zip(args, fx_node.args)
+    )
     kwargs = {k: apply_constraint(v, fx_node.kwargs[k]) for k, v in kwargs.items()}
     return args, kwargs
 
@@ -1999,21 +2001,21 @@ def tensor(data, *, dtype=None, device=None, layout=None, pin_memory=False):
     else:
         dtype = dtype or torch.get_default_dtype()
 
+    ranges: List[sympy.Expr] = []
+
     if isinstance(data, sympy.Expr):
-        ranges = []
 
         def inner_fn(index):
             return ops.index_expr(data, dtype)
 
     elif isinstance(data, (float, int)):
-        ranges = []
 
         def inner_fn(index):
             return ops.constant(data, dtype)
 
     elif len(data) == 0 or isinstance(data[0], (float, int)) and len(data) <= 8:
         # inline small tensors
-        ranges = [sympy.Integer(len(data))]
+        ranges.append(sympy.Integer(len(data)))
 
         def inner_fn(index):
             def binary_search(start, end):
@@ -2142,7 +2144,7 @@ def empty(
     assert memory_format in (None, torch.contiguous_format)
     device = decode_device(device)
     if len(size) == 1 and isinstance(size[0], (list, tuple, torch.Size)):
-        size = list(size[0])
+        size = tuple(size[0])
     return empty_strided(
         size, None, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory
     )
@@ -2184,7 +2186,7 @@ def new_constant(fill_value):
     def _new_constant(
         x, size, *, dtype=None, layout=None, device=None, pin_memory=None
     ):
-        assert isinstance(size, (list, type))
+        assert isinstance(size, (list, tuple))
         assert not pin_memory
         assert layout in (None, torch.strided)
         dtype = decode_dtype(dtype) or x.get_dtype()
@@ -2210,8 +2212,8 @@ def new_empty(x, size, *, dtype=None, layout=None, device=None, pin_memory=None)
 def empty_strided(
     size, stride, *, dtype=None, layout=None, device=None, pin_memory=None
 ):
-    assert isinstance(size, (list, type))
-    assert isinstance(stride, (list, type, type(None)))
+    assert isinstance(size, (list, tuple))
+    assert isinstance(stride, (list, tuple, type(None)))
     assert not pin_memory
     assert layout in (None, torch.strided)
     dtype = decode_dtype(dtype) or torch.get_default_dtype()
@@ -2560,7 +2562,14 @@ def scatter(x, dim: int, index, src, **kwargs):
 
 
 def scatter_fallback(
-    fn, self, dim: int, index, src, *, reduce: str = None, include_self: bool = True
+    fn,
+    self,
+    dim: int,
+    index,
+    src,
+    *,
+    reduce: Optional[str] = None,
+    include_self: bool = True,
 ):
     reduce_ty = "add" if fn == "aten.scatter_" else "sum"
     if (
@@ -2577,7 +2586,7 @@ def scatter_fallback(
 
 
 @register_lowering(aten.scatter_, type_promotion_kind=None)
-def scatter_(self, dim: int, index, src, *, reduce: str = None):
+def scatter_(self, dim: int, index, src, *, reduce: Optional[str] = None):
     assert reduce in {None, "add", "multiply"}
 
     fallback_result = scatter_fallback(
@@ -2711,7 +2720,9 @@ def backend_reduce_str(reduce):
     return self
 
 
-def upsample_nearestnd(x, output_size, scales_x: Tuple[float] = None, n: int = 2):
+def upsample_nearestnd(
+    x, output_size, scales_x: Tuple[Optional[float], ...], n: int = 2
+):
     x.realize_hint()  # elements are reused
     x_loader = x.make_loader()
     i_sizes = x.get_size()[-n:]
@@ -2726,7 +2737,7 @@ def upsample_nearestnd(x, output_size, scales_x: Tuple[float] = None, n: int = 2
         if scale:
             scales[i] = scale
 
-    def scale(x, scale, size):
+    def scale_fn(x, scale, size):
         x = ops.index_expr(x, torch.float32)
         x = ops.mul(x, ops.constant(scale, torch.float32))
         x = ops.to_dtype(x, torch.int32)
@@ -2736,7 +2747,7 @@ def fn(idx):
         x = idx[-n:]
         b = idx[:-n]
         return x_loader(
-            [*b, *[scale(i, s, size) for i, s, size in zip(x, scales, i_sizes)]]
+            [*b, *[scale_fn(i, s, size) for i, s, size in zip(x, scales, i_sizes)]]
         )
 
     return Pointwise.create(
@@ -2939,23 +2950,6 @@ def index_range_condition(index_range):
             ub = ops.index_expr(ub, torch.int64)
             return ops.and_(ops.ge(i, lb), ops.le(i, ub))
 
-        def accumulate(out_x, out_y, index_range1, index_range2=None):
-            nonlocal grad
-
-            # If the upper bound is less than the lower bound, we can get rid of one accumulation.
-            # This happens when the padding size is zero.
-            upper_less_than_lower1 = index_range1[2] < index_range1[1]
-            if isinstance(upper_less_than_lower1, bool) and upper_less_than_lower1:
-                return
-            cond = index_range_condition(index_range1)
-            if index_range2 is not None:
-                upper_less_than_lower2 = index_range2[2] < index_range2[1]
-                if isinstance(upper_less_than_lower2, bool) and upper_less_than_lower2:
-                    return
-                cond = ops.and_(cond, index_range_condition(index_range2))
-            g = ops.masked(cond, lambda: load_from_output(out_x, out_y), 0.0)
-            grad = ops.add(grad, g)
-
         # Areas after reflection:
         #
         #   top-left    |   top     |   top-right
@@ -2978,6 +2972,24 @@ def accumulate(out_x, out_y, index_range1, index_range2=None):
             index_range_condition(range_cx), index_range_condition(range_cy)
         )
         grad = ops.masked(cond, lambda: load_from_output(center_x, center_y), 0.0)
+
+        def accumulate(out_x, out_y, index_range1, index_range2=None):
+            nonlocal grad
+
+            # If the upper bound is less than the lower bound, we can get rid of one accumulation.
+            # This happens when the padding size is zero.
+            upper_less_than_lower1 = index_range1[2] < index_range1[1]
+            if isinstance(upper_less_than_lower1, bool) and upper_less_than_lower1:
+                return
+            cond = index_range_condition(index_range1)
+            if index_range2 is not None:
+                upper_less_than_lower2 = index_range2[2] < index_range2[1]
+                if isinstance(upper_less_than_lower2, bool) and upper_less_than_lower2:
+                    return
+                cond = ops.and_(cond, index_range_condition(index_range2))
+            g = ops.masked(cond, lambda: load_from_output(out_x, out_y), 0.0)
+            grad = ops.add(grad, g)
+
         accumulate(center_x, left_reflect_y, range_cx, (y, 1, left))
         accumulate(center_x, right_reflect_y, range_cx, (y, w - right, w - 1))
         accumulate(top_reflect_x, center_y, (x, 1, top), range_cy)
@@ -3076,7 +3088,7 @@ def offset_fn(index):
     )
 
 
-def range_mask_low(i: sympy.Expr, low: sympy.Expr):
+def range_mask_low(i: sympy.Expr, low: Union[sympy.Expr, int]):
     return ops.ge(
         ops.index_expr(i, torch.int64),
         ops.index_expr(sympy.Integer(low), torch.int64),
@@ -3262,8 +3274,8 @@ def max_pool2d_with_indices_backward(
         # some classes don't have `get_stride`
         # TODO will need a better way of determining if inputs are channels-last
         gO_stride = None
-    if isinstance(x, TensorBox) and isinstance(x.data.data, Pointwise):
-        data = x.data.data
+    if isinstance(x, TensorBox) and isinstance(x.data.data, Pointwise):  # type: ignore[attr-defined]
+        data = x.data.data  # type: ignore[attr-defined]
         x_buffer = ir.ComputedBuffer(
             name=None,
             layout=ir.FlexibleLayout(