mypy index propagation (pytorch#105622)

Fixes #ISSUE_NUMBER Pull Request resolved: pytorch#105622 Approved by: https://github.com/eellison

Author: msaroufim

torch/_inductor/index_propagation.py

@@ -21,7 +21,7 @@
 """
 import itertools
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Callable, Dict, Optional, Tuple, Union
 
 import sympy
 
@@ -47,11 +47,11 @@ class SymPyOps:
  """
 
  @staticmethod
- def identity(value):
+ def identity(value: Any) -> Any:
  return value
 
  @staticmethod
- def constant(value, dtype):
+ def constant(value: Union[int, float, bool], dtype: torch.dtype) -> TypedExpr:
  if is_boolean_dtype(dtype):
  expr = sympy.Integer(bool(value))
  elif is_integer_dtype(dtype):
@@ -61,13 +61,13 @@ def constant(value, dtype):
  return TypedExpr(expr, dtype)
 
  @staticmethod
- def index_expr(value, dtype):
+ def index_expr(value: sympy.Expr, dtype: torch.dtype) -> Union[int, TypedExpr]:
  if isinstance(value, int):
  value = sympy.Integer(value)
  return TypedExpr(value, dtype)
 
  @staticmethod
- def to_dtype(value, dtype):
+ def to_dtype(value: Any, dtype: torch.dtype) -> Union[int, TypedExpr]:
  if isinstance(value.expr, (sympy.Integer, sympy.Float)):
  return SymPyOps.constant(value.expr, dtype)
  elif is_integer_dtype(dtype) and is_integer_dtype(value.dtype):
@@ -77,38 +77,38 @@ def to_dtype(value, dtype):
  return NotImplemented
 
  @staticmethod
- def square(x):
+ def square(x: TypedExpr) -> TypedExpr:
  return TypedExpr(x.expr * x.expr, x.dtype)
 
  @staticmethod
- def add(x, y):
+ def add(x: TypedExpr, y: TypedExpr) -> TypedExpr:
  result_type = torch.promote_types(x.dtype, y.dtype)
  return TypedExpr(x.expr + y.expr, result_type)
 
  @staticmethod
- def sub(x, y):
+ def sub(x: TypedExpr, y: TypedExpr) -> TypedExpr:
  result_type = torch.promote_types(x.dtype, y.dtype)
  return TypedExpr(x.expr - y.expr, result_type)
 
  @staticmethod
- def mul(x, y):
+ def mul(x: TypedExpr, y: TypedExpr) -> TypedExpr:
  result_type = torch.promote_types(x.dtype, y.dtype)
  return TypedExpr(x.expr * y.expr, result_type)
 
  @staticmethod
- def neg(x):
+ def neg(x: TypedExpr) -> TypedExpr:
  return TypedExpr(-x.expr, x.dtype)
 
  @staticmethod
- def floordiv(x, y):
+ def floordiv(x: TypedExpr, y: TypedExpr) -> TypedExpr:
  result_type = torch.promote_types(x.dtype, y.dtype)
  if not is_integer_dtype(result_type):
  return NotImplemented
 
  return TypedExpr(FloorDiv(x.expr, y.expr), result_type)
 
  @staticmethod
- def remainder(x, y):
+ def remainder(x: TypedExpr, y: TypedExpr) -> Optional[TypedExpr]:
  result_type = torch.promote_types(x.dtype, y.dtype)
  if not is_integer_dtype(result_type):
  return NotImplemented
@@ -117,12 +117,12 @@ def remainder(x, y):
  return TypedExpr(result_expr, result_type)
 
  @staticmethod
- def minimum(x, y):
+ def minimum(x: TypedExpr, y: TypedExpr) -> TypedExpr:
  result_type = torch.promote_types(x.dtype, y.dtype)
  return TypedExpr(sympy.Min(x.expr, y.expr), result_type)
 
  @staticmethod
- def maximum(x, y):
+ def maximum(x: TypedExpr, y: TypedExpr) -> TypedExpr:
  result_type = torch.promote_types(x.dtype, y.dtype)
  return TypedExpr(sympy.Max(x.expr, y.expr), result_type)
 
@@ -150,18 +150,18 @@ class IndexPropagation:
 
  """
 
- def __init__(self, inner):
+ def __init__(self, inner: Any):
  self._inner = inner
 
- def materialize_expr(self, expr, dtype):
+ def materialize_expr(self, expr: sympy.Expr, dtype: torch.dtype) -> Any:
  # Construct a new constant/index_expr from the SymPy expression
  if isinstance(expr, sympy.Integer):
  return self._inner.constant(int(expr), dtype)
  elif not expr.free_symbols:
  return self._inner.constant(float(expr), dtype)
  return self._inner.index_expr(expr, dtype)
 
- def unwrap(self, a):
+ def unwrap(self, a: Union[Any, IndexPropVar]) -> Any:
  if not isinstance(a, IndexPropVar):
  return a
 
@@ -171,15 +171,17 @@ def unwrap(self, a):
 
  return a.value
 
- def fallback(self, name, args, kwargs):
+ def fallback(self, name: str, args: Tuple, kwargs: Dict[str, Any]) -> IndexPropVar:
  # Fallback to the wrapped handler
  new_args = [self.unwrap(a) for a in args]
  new_kwargs = {k: self.unwrap(v) for k, v in kwargs.items()}
  return IndexPropVar(getattr(self._inner, name)(*new_args, **new_kwargs))
 
- def propagate_sympy(self, name, args, kwargs):
+ def propagate_sympy(
+ self, name: str, args: Tuple, kwargs: Dict[str, Any]
+ ) -> Union[TypedExpr, IndexPropVar]:
  # Build a new SymPy expression from this ops call
- def unwrap(a):
+ def unwrap(a: Union[Any, IndexPropVar]) -> Any:
  if not isinstance(a, IndexPropVar):
  return a
  return a.value
@@ -191,8 +193,8 @@ def unwrap(a):
  return self.fallback(name, args, kwargs)
  return IndexPropVar.new_symbolic(new_expr)
 
- def __getattr__(self, name):
- def inner(*args, **kwargs):
+ def __getattr__(self, name: str) -> Callable[..., Union[Any, IndexPropVar]]:
+ def inner(*args: Any, **kwargs: Any) -> Union[Any, IndexPropVar]:
  if not hasattr(SymPyOps, name):
  return self.fallback(name, args, kwargs)
 
@@ -208,7 +210,9 @@ def inner(*args, **kwargs):
 
  return inner
 
- def indirect_indexing(self, index, size, check=True):
+ def indirect_indexing(
+ self, index: Union[Any, IndexPropVar], size: Any, check: bool = True
+ ) -> Any:
  # indirect_indexing returns a sympy value, so no need to wrap in IndexPropVar here
  if isinstance(index, IndexPropVar) and index.is_symbolic:
  return index.value.expr