Add type information to Relay's Python AST (#8)

* begins typing expr.py

* minor style fix

* improves typing and adds more

* type -> typ to avoid keyword clash

* reverts prev change

* fixes typo

* merge python/expr conflicts

* incorporate changes

* lint test

* comment out instance var lint test

* undo tests

* adds more types

* more types

* disable specific pylint invalid-name warnings

* minor comments on particular nodes

* minor clean up

relay/python/relay/expr.py

@@ -1,7 +1,7 @@
 # pylint: disable=no-else-return, unidiomatic-typecheck
 """All the expression nodes"""
 from enum import IntEnum
-from typing import Union
+from typing import Union, Tuple, Dict
 from .base import register_nnvm_node, NodeBase
 from . import make
 
@@ -13,6 +13,8 @@ class Environment(NodeBase):
  """Global Environment
  """
 
+ items: Dict["GlobalId", "Item"]
+
  def add(self, func: "GlobalId") -> None:
  return make.Environment_add(self, func)
 
@@ -33,6 +35,9 @@ class Item(NodeBase):
  """Base class of all expressions.
  """
 
+ #pylint: disable=invalid-name
+ id: "GlobalId"
+
  def __eq__(self, other):
  if type(self) != type(other):
  return False
@@ -53,11 +58,13 @@ class Type(NodeBase):
  def __eq__(self, other):
  return alpha_eq(self, other)
 
+
 @register_nnvm_node
 class Attributes(NodeBase):
  def __getitem__(self, index):
  return self.attributes[index]
 
+
 class Builder(object):
  """A helper class for building partial AST fragments."""
 
@@ -128,88 +135,97 @@ def __eq__(self, other):
 
 @register_nnvm_node
 class Defn(Item):
- pass
+ type: Type
+ body: Expr
 
 
 @register_nnvm_node
 class Primitive(Item):
- pass
+ type: Type
 
 
 @register_nnvm_node
 class String(Expr):
- pass
+ value: str
 
 
 @register_nnvm_node
 class IntLit(Expr):
- pass
+ value: int
 
 
 @register_nnvm_node
 class FloatLit(Expr):
- pass
+ value: float
 
 
 @register_nnvm_node
 class BoolLit(Expr):
- pass
+ value: bool
 
 
 @register_nnvm_node
 class TensorLit(Expr):
- pass
+ data: Tuple[NodeBase, ...]
 
 
 @register_nnvm_node
 class ProductLit(Expr):
- pass
+ fields: Tuple[NodeBase, ...]
 
 
 @register_nnvm_node
 class BaseType(Type):
- pass
+ type: str
 
 
 @register_nnvm_node
 class Cast(Expr):
- pass
+ target: Type
+ node: Expr
 
 
 @register_nnvm_node
 class LocalId(Expr):
+ name: str
+
  def __hash__(self) -> int:
  return hash(self.name)
 
 
 @register_nnvm_node
 class GlobalId(Expr):
- pass
+ name: str
 
 
 @register_nnvm_node
 class IntrinsicId(Expr):
- pass
+ name: str
 
 
 @register_nnvm_node
 class Param(Expr):
- pass
+ #pylint: disable=invalid-name
+ id: LocalId
+ type: Type
 
 
 @register_nnvm_node
 class Function(Expr):
- pass
+ params: Tuple[Param, ...]
+ body: Expr
 
 
 @register_nnvm_node
 class Call(Expr):
- pass
+ #pylint: disable=invalid-name
+ fn: Expr
+ args: Tuple[Expr, ...]
 
 
 @register_nnvm_node
 class Debug(Expr):
- pass
+ node: Expr
 
 
 class UOp(IntEnum):
@@ -218,11 +234,11 @@ class UOp(IntEnum):
 
 @register_nnvm_node
 class UnaryOp(Expr):
- pass
-
-#pylint: disable=invalid-name
+ op: UOp
+ node: Expr
 
 
+#pylint: disable=invalid-name
 class BOp(IntEnum):
  """The set of builtin binary ops supported by Relay."""
  PLUS = 0
@@ -239,12 +255,16 @@ class BOp(IntEnum):
 
 @register_nnvm_node
 class BinaryOp(Expr):
- pass
+ op: BOp
+ left: Expr
+ right: Expr
 
 
 @register_nnvm_node
 class Let(Expr):
- pass
+ id: LocalId
+ value: Expr
+ body: Expr
 
 
 @register_nnvm_node
@@ -254,51 +274,61 @@ class Functor(Expr):
 
 @register_nnvm_node
 class Reverse(Expr):
- pass
+ node: Expr
 
 
 @register_nnvm_node
 class Accumulate(Expr):
- pass
+ update_binders: Tuple[LocalId, ...]
+ value: NodeBase
 
 
 @register_nnvm_node
 class Zero(Expr):
- pass
+ type: Type
 
-# TODO(@jroesch): todo move
 
+# TODO(@jroesch): todo move
+# values
 
 @register_nnvm_node
 class IntValue(Value):
- pass
+ value: int
 
 
 @register_nnvm_node
 class FloatValue(Value):
- pass
+ value: float
 
 
 @register_nnvm_node
 class BoolValue(Value):
- pass
+ value: bool
 
 
 @register_nnvm_node
 class FnValue(Value):
- pass
+ env: Dict[LocalId, Value]
+ func: Function
+
+# unsorted
 
 @register_nnvm_node
 class If(Expr):
- pass
+ guard: Expr
+ true_b: Expr
+ false_b: Expr
+
 
 @register_nnvm_node
 class Shape(NodeBase):
  pass
 
+
 @register_nnvm_node
 class TensorType(Type):
- pass
+ dtype: BaseType
+ shape: Shape
 
 
 def alpha_eq(left: Union[Expr, Type], right: Union[Expr, Type]) -> bool:

relay/python/relay/relay.py

@@ -17,9 +17,6 @@ def ensure(expected, value, msg):
  if value != expected:
  raise Exception("{}: found {}".format(msg, value))
 
-# TODO: it would be great to report errors in the Python program
-
-
 def compile_args_to_params(args):
  """This function will convert Python arguments into Relay parameters."""
  ensure([], args.defaults, "relay decorator does not support default arguments")
@@ -87,7 +84,6 @@ def relay_type_from_annotation(annotation):
 #
 # Process a single definition and produce a single Relay Defunc.
 
-
 class DefToRelay(ast.NodeVisitor):
  """Compiles a Python definition to a Relay definition."""
  # local_scopes: List[Dict[LocalId, Expr]]
@@ -242,8 +238,6 @@ def wrapper(*_):
 
 # Store Python line and columb information for errors
 #
-# Make it possible to handle calls with keywords arguments
-#
 # Ensure names prefixed with relay. becomes a Intrinsic
 #
 # Handle assignments