From c9b82e3903b0837dbbce1dc2c211a58b9b098321 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 13 Apr 2020 15:10:48 -0700
Subject: [PATCH 01/10] Cleanup type pack and unpack for tuples.

---
 include/tvm/relay/attrs/device_copy.h |   4 +-
 include/tvm/relay/attrs/memory.h      |  24 ++++++
 python/tvm/relay/op/memory/memory.py  |  16 +++-
 src/relay/op/memory/memory.cc         | 103 ++++++++++++++++++++++----
 4 files changed, 130 insertions(+), 17 deletions(-)

diff --git a/include/tvm/relay/attrs/device_copy.h b/include/tvm/relay/attrs/device_copy.h
index 393562980308..2486fcdf473d 100644
--- a/include/tvm/relay/attrs/device_copy.h
+++ b/include/tvm/relay/attrs/device_copy.h
@@ -40,11 +40,11 @@ struct DeviceCopyAttrs : public tvm::AttrsNode<DeviceCopyAttrs> {
   TVM_DECLARE_ATTRS(DeviceCopyAttrs, "relay.attrs.DeviceCopyAttrs") {
     TVM_ATTR_FIELD(src_dev_type)
       .describe(
-         "The virutal device/context type where the op copies data from.")
+         "The virtual device/context type where the op copies data from.")
       .set_default(0);
     TVM_ATTR_FIELD(dst_dev_type)
       .describe(
-         "The virutal device/context type where the op copies data to.")
+         "The virtual device/context type where the op copies data to.")
       .set_default(0);
   }
 };
diff --git a/include/tvm/relay/attrs/memory.h b/include/tvm/relay/attrs/memory.h
index 00204b315a1b..41e0f036544b 100644
--- a/include/tvm/relay/attrs/memory.h
+++ b/include/tvm/relay/attrs/memory.h
@@ -31,6 +31,30 @@
 namespace tvm {
 namespace relay {
 
+/*!
+ * \brief Options for allocating storage.
+ */
+struct AllocStorageAttrs : public tvm::AttrsNode<AllocStorageAttrs> {
+  DataType dtype;
+  int device_id;
+  int device_type;
+
+  TVM_DECLARE_ATTRS(AllocStorageAttrs, "relay.attrs.AllocStorageAttrs") {
+    TVM_ATTR_FIELD(dtype)
+      .describe(
+         "The dtype of the tensor to allocate.")
+      .set_default(DataType::Float(32, 1));
+    TVM_ATTR_FIELD(device_id)
+      .describe(
+        "The device id on which to allocate memory."
+      );
+    TVM_ATTR_FIELD(device_type)
+      .describe(
+        "The device type on which to allocate memory."
+      );
+  }
+};
+
 /*!
  * \brief Options for allocating tensors.
  */
diff --git a/python/tvm/relay/op/memory/memory.py b/python/tvm/relay/op/memory/memory.py
index 892ba88c17c5..f43cd8b407b3 100644
--- a/python/tvm/relay/op/memory/memory.py
+++ b/python/tvm/relay/op/memory/memory.py
@@ -23,6 +23,9 @@ def invoke_tvm_op(func, inputs, outputs):
 
     Parameters
     ----------
+    func : tvm.relay.Expr
+        The input expr.
+
     inputs : tvm.relay.Expr
         A tuple of the inputs to pass to the TVM function.
 
@@ -59,7 +62,7 @@ def alloc_tensor(storage, shape, dtype='float32', assert_shape=None):
     """
     return _make.alloc_tensor(storage, shape, dtype, assert_shape)
 
-def alloc_storage(size, alignment, dtype_hint='float32'):
+def alloc_storage(size, alignment, ctx, dtype_hint='float32'):
     """Allocate a piece of tensor storage.
 
     Parameters
@@ -76,7 +79,7 @@ def alloc_storage(size, alignment, dtype_hint='float32'):
     result : tvm.relay.Expr
         The alloc_storage expression.
     """
-    return _make.alloc_storage(size, alignment, dtype_hint)
+    return _make.alloc_storage(size, alignment, dtype_hint, ctx)
 
 def shape_func(func, inputs, outputs, dependent=False):
     """Invoke the shape function of the passed function.
@@ -96,3 +99,12 @@ def shape_func(func, inputs, outputs, dependent=False):
         The shape function expression.
     """
     return _make.shape_func(func, inputs, outputs, dependent)
+
+def flatten_tuple_type(ty):
+    return _make.FlattenTupleType(ty)
+
+def from_tuple_type(ty, expr):
+    return _make.FromTupleType(ty, expr)
+
+def to_tuple_type(ty, exprs):
+    return _make.ToTupleType(ty, exprs)
diff --git a/src/relay/op/memory/memory.cc b/src/relay/op/memory/memory.cc
index c9ab067da594..4a85dfe0a15e 100644
--- a/src/relay/op/memory/memory.cc
+++ b/src/relay/op/memory/memory.cc
@@ -35,6 +35,7 @@
 namespace tvm {
 namespace relay {
 
+TVM_REGISTER_NODE_TYPE(AllocStorageAttrs);
 TVM_REGISTER_NODE_TYPE(AllocTensorAttrs);
 TVM_REGISTER_NODE_TYPE(ShapeFuncAttrs);
 
@@ -42,9 +43,11 @@ TVM_REGISTER_NODE_TYPE(ShapeFuncAttrs);
 // We should consider a better solution, i.e the type relation
 // being able to see the arguments as well?
 TVM_REGISTER_GLOBAL("relay.op.memory._make.alloc_storage")
-    .set_body_typed([](Expr size, Expr alignment, DataType dtype) {
-      auto attrs = make_object<AllocTensorAttrs>();
+    .set_body_typed([](Expr size, Expr alignment, DataType dtype, TVMContext ctx) {
+      auto attrs = make_object<AllocStorageAttrs>();
       attrs->dtype = dtype;
+      attrs->device_id = ctx.device_id;
+      attrs->device_type = ctx.device_type;
       static const Op& op = Op::Get("memory.alloc_storage");
       return Call(op, {size, alignment}, Attrs(attrs), {});
     });
@@ -209,10 +212,20 @@ bool InvokeTVMOPRel(const Array<Type>& types, int num_inputs, const Attrs& attrs
 }
 
 TVM_REGISTER_GLOBAL("relay.op.memory._make.invoke_tvm_op")
-    .set_body_typed(
-        [](Expr func, Expr inputs, Expr outputs) {
-          return Call(Op::Get("memory.invoke_tvm_op"), {func, inputs, outputs}, Attrs());
-        });
+.set_body_typed(
+  [](Expr func, Expr inputs, Expr outputs) {
+    Attrs attrs;
+    // Record the attribute of the input expression. The attribute of the master
+    // op in a fused function is used.
+    if (const auto* fn = func.as<FunctionNode>()) {
+      if (const auto* cn = fn->body.as<CallNode>()) {
+        attrs = cn->attrs;
+      }
+    } else if (const auto* cn = func.as<CallNode>()) {
+      attrs = cn->attrs;
+    }
+    return Call(Op::Get("memory.invoke_tvm_op"), {func, inputs, outputs}, attrs);
+  });
 
 RELAY_REGISTER_OP("memory.invoke_tvm_op")
     .describe(R"code(Invoke an operation compiled by TVM.)code" TVM_ADD_FILELINE)
@@ -265,29 +278,93 @@ TVM_REGISTER_GLOBAL("relay.op.memory._make.shape_func")
       return Call(op, {func, inputs, outputs}, Attrs(attrs), {});
     });
 
-static void FlattenTypeAux(const Type& type, std::vector<TensorType>* out) {
+// # TODO(@jroesch): port to c++ and unify with existing code
+// class LinearizeRetType:
+//     """A linear view of a Relay type, handles a linear order
+//        for nested tuples, and tensor types.
+//     """
+
+//     def __init__(self, typ):
+//         """Initialize the linearizer."""
+//         self.typ = typ
+
+//     def unpack(self):
+//         """Return the linear representation of the type."""
+//         def _unpack(typ, out):
+//             # TODO(@jroesch): replace with new flattening pass
+//             if isinstance(typ, ty.TensorType):
+//                 out.append(typ)
+//             elif isinstance(typ, ty.TupleType):
+//                 for field_ty in typ.fields:
+//                     _unpack(field_ty, out)
+//             else:
+//                 raise Exception("unsupported Relay type: {0}".format(typ))
+
+//         output = []
+//         _unpack(self.typ, output)
+//         return output
+
+//     def pack(self, seq):
+//         """Repack a linear type as a nested type."""
+//         def _pack(value, typ, out):
+//             if isinstance(typ, ty.TensorType):
+//                 out.append(value)
+//             elif isinstance(typ, ty.TupleType):
+//                 tuple_out = []
+//                 for i, field_ty in enumerate(typ.fields):
+//                     _pack(value[i], field_ty, tuple_out)
+//                 out.append(expr.Tuple(tuple_out))
+//             else:
+//                 raise Exception("unsupported Relay type: {0}".format(typ))
+
+//         if len(seq) == 1:
+//             return seq[0]
+//         else:
+//             out = []
+//             _pack(seq, self.typ, out)
+//             assert len(out) == 1, "must return fully packed type"
+//             return out[0]
+
+static void FlattenTupleTypeAux(const Type& type, std::vector<TensorType>* out) {
   if (auto tt = type.as<TensorTypeNode>()) {
     out->push_back(GetRef<TensorType>(tt));
   } else if (auto tuple_ty = type.as<TupleTypeNode>()) {
     for (auto field : tuple_ty->fields) {
-      FlattenTypeAux(field, out);
+      FlattenTupleTypeAux(field, out);
     }
   } else {
     LOG(FATAL) << "unsupported " << type;
   }
 }
 
-std::vector<TensorType> FlattenType(const Type& type) {
+std::vector<TensorType> FlattenTupleType(const Type& type) {
   std::vector<TensorType> out;
-  FlattenTypeAux(type, &out);
+  FlattenTupleTypeAux(type, &out);
   return out;
 }
 
-Expr PackByType(const Type& t, const Array<Expr>& exprs) {
+Array<Expr> FromTupleType(const Type& type, const Expr& expr) {
+  LOG(FATAL) << "NYI";
+}
+
+// Pack the sequence of expressions according to the provided TupleType.
+Expr ToTupleType(const Type& t, const Array<Expr>& exprs) {
   LOG(FATAL) << "NYI";
   return Expr();
 }
 
+TVM_REGISTER_GLOBAL("relay.op.memory._make.FlattenTupleType")
+.set_body_typed([](Type type) {
+  auto types = FlattenTupleType(type);
+  return Array<Type>(types.begin(), types.end());
+});
+
+TVM_REGISTER_GLOBAL("relay.op.memory._make.FromTupleType")
+.set_body_typed(FromTupleType);
+
+TVM_REGISTER_GLOBAL("relay.op.memory._make.ToTupleType")
+.set_body_typed(ToTupleType);
+
 bool ShapeFuncRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
                   const TypeReporter& reporter) {
   CHECK_EQ(types.size(), 4u);
@@ -298,8 +375,8 @@ bool ShapeFuncRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
   CHECK(func_type != nullptr);
 
   auto tuple = TupleType(func_type->arg_types);
-  auto in_types = FlattenType(tuple);
-  auto out_types = FlattenType(func_type->ret_type);
+  auto in_types = FlattenTupleType(tuple);
+  auto out_types = FlattenTupleType(func_type->ret_type);
 
   Array<Type> shape_func_ins, shape_func_outs;
   for (size_t i = 0; i < in_types.size(); i++) {

From 6354e8e19fd1408314693fe03cccd05474586280 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 13 Apr 2020 16:47:07 -0700
Subject: [PATCH 02/10] Clean up the memory_pass using common helpers

---
 include/tvm/relay/attrs/memory.h           |  4 +
 python/tvm/relay/op/memory/memory.py       |  2 +-
 python/tvm/relay/transform/memory_alloc.py | 79 +++--------------
 src/relay/backend/vm/compiler.cc           |  2 +-
 src/relay/op/memory/memory.cc              | 99 ++++++++++++----------
 5 files changed, 71 insertions(+), 115 deletions(-)

diff --git a/include/tvm/relay/attrs/memory.h b/include/tvm/relay/attrs/memory.h
index 41e0f036544b..db4e04eba54f 100644
--- a/include/tvm/relay/attrs/memory.h
+++ b/include/tvm/relay/attrs/memory.h
@@ -31,6 +31,10 @@
 namespace tvm {
 namespace relay {
 
+std::vector<TensorType> FlattenTupleType(const Type& type);
+std::vector<Expr> FromTupleType(const Type& type, const Expr& expr);
+Expr ToTupleType(const Type& t, const Array<Expr>& exprs);
+
 /*!
  * \brief Options for allocating storage.
  */
diff --git a/python/tvm/relay/op/memory/memory.py b/python/tvm/relay/op/memory/memory.py
index f43cd8b407b3..b3488c6d93dd 100644
--- a/python/tvm/relay/op/memory/memory.py
+++ b/python/tvm/relay/op/memory/memory.py
@@ -79,7 +79,7 @@ def alloc_storage(size, alignment, ctx, dtype_hint='float32'):
     result : tvm.relay.Expr
         The alloc_storage expression.
     """
-    return _make.alloc_storage(size, alignment, dtype_hint, ctx)
+    return _make.alloc_storage(size, alignment, ctx, dtype_hint)
 
 def shape_func(func, inputs, outputs, dependent=False):
     """Invoke the shape function of the passed function.
diff --git a/python/tvm/relay/transform/memory_alloc.py b/python/tvm/relay/transform/memory_alloc.py
index c238730807d3..c0e040d93899 100644
--- a/python/tvm/relay/transform/memory_alloc.py
+++ b/python/tvm/relay/transform/memory_alloc.py
@@ -26,60 +26,14 @@
 from ... import DataType, register_func
 from .. import ty, expr
 from ..backend import compile_engine
+from ..op.memory import flatten_tuple_type, from_tuple_type, to_tuple_type
+from ...import cpu
 
 
 def is_primitive(call):
     return hasattr(call, 'op') and hasattr(call.op, 'attrs') and \
            hasattr(call.op.attrs, 'Primitive') and int(call.op.attrs.Primitive) == 1
 
-# TODO(@jroesch): port to c++ and unify with existing code
-class LinearizeRetType:
-    """A linear view of a Relay type, handles a linear order
-       for nested tuples, and tensor types.
-    """
-
-    def __init__(self, typ):
-        """Initialize the linearizer."""
-        self.typ = typ
-
-    def unpack(self):
-        """Return the linear representation of the type."""
-        def _unpack(typ, out):
-            # TODO(@jroesch): replace with new flattening pass
-            if isinstance(typ, ty.TensorType):
-                out.append(typ)
-            elif isinstance(typ, ty.TupleType):
-                for field_ty in typ.fields:
-                    _unpack(field_ty, out)
-            else:
-                raise Exception("unsupported Relay type: {0}".format(typ))
-
-        output = []
-        _unpack(self.typ, output)
-        return output
-
-    def pack(self, seq):
-        """Repack a linear type as a nested type."""
-        def _pack(value, typ, out):
-            if isinstance(typ, ty.TensorType):
-                out.append(value)
-            elif isinstance(typ, ty.TupleType):
-                tuple_out = []
-                for i, field_ty in enumerate(typ.fields):
-                    _pack(value[i], field_ty, tuple_out)
-                out.append(expr.Tuple(tuple_out))
-            else:
-                raise Exception("unsupported Relay type: {0}".format(typ))
-
-        if len(seq) == 1:
-            return seq[0]
-        else:
-            out = []
-            _pack(seq, self.typ, out)
-            assert len(out) == 1, "must return fully packed type"
-            return out[0]
-
-
 class ManifestAllocPass(ExprMutator):
     """A pass for explictly manifesting all memory allocations in Relay."""
 
@@ -90,6 +44,7 @@ def __init__(self, target_host):
         self.shape_func = op.memory.shape_func
         self.scopes = [ScopeBuilder()]
         self.target_host = target_host
+        self.default_context = cpu(0)
         self.compute_dtype = "int64"
         super().__init__()
 
@@ -147,7 +102,7 @@ def make_static_allocation(self, scope, tensor_type, i):
         alignment = self.compute_alignment(tensor_type.dtype)
         dtype = tensor_type.dtype
         sto = scope.let("storage_{0}".format(i), self.alloc_storage(
-            size, alignment, dtype))
+            size, alignment, self.default_context, dtype))
         # TODO(@jroesch): There is a bug with typing based on the constant shape.
         tensor = self.alloc_tensor(sto, shape, dtype, tensor_type.shape)
         return scope.let("tensor_{0}".format(i), tensor)
@@ -174,8 +129,7 @@ def visit_call(self, call):
             new_args = [self.visit(arg) for arg in call.args]
             ins = expr.Tuple(new_args)
             ret_type = call.checked_type
-            view = LinearizeRetType(ret_type)
-            out_types = view.unpack()
+            out_types = flatten_tuple_type(ret_type)
 
             is_dynamic = ty.type_has_any(ret_type)
             # TODO(@jroesch): restore this code, more complex then it seems
@@ -194,18 +148,11 @@ def visit_call(self, call):
                     state = int(state)
                     # Pass Shapes
                     if state == 2:
-                        if isinstance(arg.type_annotation, ty.TupleType):
-                            for j in range(len(arg.type_annotation.fields)):
-                                let_in_arg = scope.let("in_arg_{0}".format(input_pos + j),
-                                                       expr.TupleGetItem(arg, j))
-                                sh_of = self.visit(self.shape_of(let_in_arg))
-                                shape_func_ins.append(
-                                    scope.let("in_shape_{0}".format(input_pos + j), sh_of))
-                            input_pos += len(arg.type_annotation.fields)
-                        else:
-                            sh_of = self.visit(self.shape_of(arg))
+                        for j, subexp in enumerate(from_tuple_type(arg.type_annotation, arg)):
+                            let_in_arg = scope.let("in_arg_{0}".format(input_pos + j), subexp)
+                            sh_of = self.visit(self.shape_of(let_in_arg))
                             shape_func_ins.append(
-                                scope.let("in_shape_{0}".format(input_pos), sh_of))
+                                scope.let("in_shape_{0}".format(input_pos + j), sh_of))
                             input_pos += 1
                         is_inputs.append(0)
                     # Pass Inputs
@@ -215,8 +162,8 @@ def visit_call(self, call):
                             scope.let("in_shape_{0}".format(input_pos), new_arg))
                         input_pos += 1
                         is_inputs.append(1)
-                    # TODO(@jroesch): handle 3rd case
                     else:
+                        # TODO(@jroesch): handle 3rd case
                         raise Exception("unsupported shape function input state")
 
                 out_shapes = []
@@ -239,7 +186,7 @@ def visit_call(self, call):
                         out_shape, out_type.dtype)
                     alignment = self.compute_alignment(out_type.dtype)
                     sto = scope.let("storage_{i}".format(i=i), self.alloc_storage(
-                        size, alignment, out_type.dtype))
+                        size, alignment, self.default_context, out_type.dtype))
                     storages.append(sto)
 
                 outs = []
@@ -256,7 +203,7 @@ def visit_call(self, call):
                 tuple_outs = expr.Tuple(outs)
                 invoke = self.invoke_tvm(call.op, ins, tuple_outs)
                 scope.let("", invoke)
-                return outs[0] if len(outs) == 1 else tuple_outs
+                return to_tuple_type(ret_type, tuple_outs.fields)
             else:
                 outs = []
                 for i, out_ty in enumerate(out_types):
@@ -266,7 +213,7 @@ def visit_call(self, call):
                 output = expr.Tuple(outs)
                 invoke = self.invoke_tvm(call.op, ins, output)
                 scope.let("", invoke)
-                return view.pack(output)
+                return to_tuple_type(ret_type, output.fields)
         else:
             return super().visit_call(call)
 
diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
index 8af6247fc810..7e2d43e7b35d 100644
--- a/src/relay/backend/vm/compiler.cc
+++ b/src/relay/backend/vm/compiler.cc
@@ -579,7 +579,7 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
           auto alignment_register = last_register_;
 
           // Get the dtype hint from the attributes.
-          auto alloc_attrs = attrs.as<AllocTensorAttrs>();
+          auto alloc_attrs = attrs.as<AllocStorageAttrs>();
           CHECK(alloc_attrs != nullptr)
               << "must be the alloc tensor attrs";
           auto dtype = alloc_attrs->dtype;
diff --git a/src/relay/op/memory/memory.cc b/src/relay/op/memory/memory.cc
index 4a85dfe0a15e..1e923ec05975 100644
--- a/src/relay/op/memory/memory.cc
+++ b/src/relay/op/memory/memory.cc
@@ -43,9 +43,9 @@ TVM_REGISTER_NODE_TYPE(ShapeFuncAttrs);
 // We should consider a better solution, i.e the type relation
 // being able to see the arguments as well?
 TVM_REGISTER_GLOBAL("relay.op.memory._make.alloc_storage")
-    .set_body_typed([](Expr size, Expr alignment, DataType dtype, TVMContext ctx) {
+    .set_body_typed([](Expr size, Expr alignment, TVMContext ctx, DataType dtype_hint) {
       auto attrs = make_object<AllocStorageAttrs>();
-      attrs->dtype = dtype;
+      attrs->dtype = dtype_hint;
       attrs->device_id = ctx.device_id;
       attrs->device_type = ctx.device_type;
       static const Op& op = Op::Get("memory.alloc_storage");
@@ -213,19 +213,9 @@ bool InvokeTVMOPRel(const Array<Type>& types, int num_inputs, const Attrs& attrs
 
 TVM_REGISTER_GLOBAL("relay.op.memory._make.invoke_tvm_op")
 .set_body_typed(
-  [](Expr func, Expr inputs, Expr outputs) {
-    Attrs attrs;
-    // Record the attribute of the input expression. The attribute of the master
-    // op in a fused function is used.
-    if (const auto* fn = func.as<FunctionNode>()) {
-      if (const auto* cn = fn->body.as<CallNode>()) {
-        attrs = cn->attrs;
-      }
-    } else if (const auto* cn = func.as<CallNode>()) {
-      attrs = cn->attrs;
-    }
-    return Call(Op::Get("memory.invoke_tvm_op"), {func, inputs, outputs}, attrs);
-  });
+        [](Expr func, Expr inputs, Expr outputs) {
+          return Call(Op::Get("memory.invoke_tvm_op"), {func, inputs, outputs}, Attrs());
+        });
 
 RELAY_REGISTER_OP("memory.invoke_tvm_op")
     .describe(R"code(Invoke an operation compiled by TVM.)code" TVM_ADD_FILELINE)
@@ -278,31 +268,6 @@ TVM_REGISTER_GLOBAL("relay.op.memory._make.shape_func")
       return Call(op, {func, inputs, outputs}, Attrs(attrs), {});
     });
 
-// # TODO(@jroesch): port to c++ and unify with existing code
-// class LinearizeRetType:
-//     """A linear view of a Relay type, handles a linear order
-//        for nested tuples, and tensor types.
-//     """
-
-//     def __init__(self, typ):
-//         """Initialize the linearizer."""
-//         self.typ = typ
-
-//     def unpack(self):
-//         """Return the linear representation of the type."""
-//         def _unpack(typ, out):
-//             # TODO(@jroesch): replace with new flattening pass
-//             if isinstance(typ, ty.TensorType):
-//                 out.append(typ)
-//             elif isinstance(typ, ty.TupleType):
-//                 for field_ty in typ.fields:
-//                     _unpack(field_ty, out)
-//             else:
-//                 raise Exception("unsupported Relay type: {0}".format(typ))
-
-//         output = []
-//         _unpack(self.typ, output)
-//         return output
 
 //     def pack(self, seq):
 //         """Repack a linear type as a nested type."""
@@ -343,14 +308,49 @@ std::vector<TensorType> FlattenTupleType(const Type& type) {
   return out;
 }
 
-Array<Expr> FromTupleType(const Type& type, const Expr& expr) {
-  LOG(FATAL) << "NYI";
+static void FromTupleTypeAux(const Type& type, const Expr& expr, std::vector<Expr>* out) {
+  if (auto tt = type.as<TensorTypeNode>()) {
+    out->push_back(expr);
+  } else if (auto tuple_ty = type.as<TupleTypeNode>()) {
+    for (auto i = 0; i < tuple_ty->fields.size(); i++) {
+      FromTupleTypeAux(tuple_ty->fields[i], TupleGetItem(expr, i), out);
+    }
+  } else {
+    LOG(FATAL) << "unsupported " << type;
+  }
+}
+
+std::vector<Expr> FromTupleType(const Type& type, const Expr& expr) {
+  std::vector<Expr> out;
+  FromTupleTypeAux(type, expr, &out);
+  return out;
+}
+
+static void ToTupleTypeAux(const Type& type, const std::vector<Expr>& exprs, int* index, std::vector<Expr>* out) {
+  if (auto tt = type.as<TensorTypeNode>()) {
+    out->push_back(exprs[*index]);
+    *index += 1;
+  } else if (auto tuple_ty = type.as<TupleTypeNode>()) {
+    std::vector<Expr> tuple_out;
+    for (auto i = 0; i < tuple_ty->fields.size(); i++) {
+      ToTupleTypeAux(tuple_ty->fields[i], exprs, index, &tuple_out);
+    }
+    out->push_back(Tuple(tuple_out));
+  } else {
+    LOG(FATAL) << "unsupported " << type;
+  }
 }
 
 // Pack the sequence of expressions according to the provided TupleType.
-Expr ToTupleType(const Type& t, const Array<Expr>& exprs) {
-  LOG(FATAL) << "NYI";
-  return Expr();
+Expr ToTupleType(const Type& t, const std::vector<Expr>& exprs) {
+  if (t.as<TensorTypeNode>() && exprs.size() == 1) {
+    return exprs[0];
+  } else {
+    std::vector<Expr> out;
+    int index = 0;
+    ToTupleTypeAux(t, exprs, &index, &out);
+    return out[0];
+  }
 }
 
 TVM_REGISTER_GLOBAL("relay.op.memory._make.FlattenTupleType")
@@ -360,10 +360,15 @@ TVM_REGISTER_GLOBAL("relay.op.memory._make.FlattenTupleType")
 });
 
 TVM_REGISTER_GLOBAL("relay.op.memory._make.FromTupleType")
-.set_body_typed(FromTupleType);
+.set_body_typed([](Type type, Expr expr) {
+  auto exprs = FromTupleType(type, expr);
+  return Array<Expr>(exprs.begin(), exprs.end());
+});
 
 TVM_REGISTER_GLOBAL("relay.op.memory._make.ToTupleType")
-.set_body_typed(ToTupleType);
+.set_body_typed([](Type t, Array<Expr> array) {
+  return ToTupleType(t, std::vector<Expr>(array.begin(), array.end()));
+});
 
 bool ShapeFuncRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
                   const TypeReporter& reporter) {

From 1614aa021aebc176173d3dedcfa175fb7630e32d Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 13 Apr 2020 16:48:10 -0700
Subject: [PATCH 03/10] Clean up memory.cc

---
 src/relay/op/memory/memory.cc | 22 ----------------------
 1 file changed, 22 deletions(-)

diff --git a/src/relay/op/memory/memory.cc b/src/relay/op/memory/memory.cc
index 1e923ec05975..54ff3bf29704 100644
--- a/src/relay/op/memory/memory.cc
+++ b/src/relay/op/memory/memory.cc
@@ -268,28 +268,6 @@ TVM_REGISTER_GLOBAL("relay.op.memory._make.shape_func")
       return Call(op, {func, inputs, outputs}, Attrs(attrs), {});
     });
 
-
-//     def pack(self, seq):
-//         """Repack a linear type as a nested type."""
-//         def _pack(value, typ, out):
-//             if isinstance(typ, ty.TensorType):
-//                 out.append(value)
-//             elif isinstance(typ, ty.TupleType):
-//                 tuple_out = []
-//                 for i, field_ty in enumerate(typ.fields):
-//                     _pack(value[i], field_ty, tuple_out)
-//                 out.append(expr.Tuple(tuple_out))
-//             else:
-//                 raise Exception("unsupported Relay type: {0}".format(typ))
-
-//         if len(seq) == 1:
-//             return seq[0]
-//         else:
-//             out = []
-//             _pack(seq, self.typ, out)
-//             assert len(out) == 1, "must return fully packed type"
-//             return out[0]
-
 static void FlattenTupleTypeAux(const Type& type, std::vector<TensorType>* out) {
   if (auto tt = type.as<TensorTypeNode>()) {
     out->push_back(GetRef<TensorType>(tt));

From 7bbffa7fd9a4923a987e68323c72c6774b0e1979 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 13 Apr 2020 16:55:44 -0700
Subject: [PATCH 04/10] Refactor pass

---
 python/tvm/relay/transform/memory_alloc.py | 153 +++++++++++----------
 1 file changed, 80 insertions(+), 73 deletions(-)

diff --git a/python/tvm/relay/transform/memory_alloc.py b/python/tvm/relay/transform/memory_alloc.py
index c0e040d93899..db5df4181964 100644
--- a/python/tvm/relay/transform/memory_alloc.py
+++ b/python/tvm/relay/transform/memory_alloc.py
@@ -122,6 +122,82 @@ def visit_let(self, let):
 
         return scope.get()
 
+    def dynamic_invoke(self, scope, op, ins, new_args, out_types, ret_type):
+        shape_func_ins = []
+        engine = compile_engine.get()
+        cfunc = engine.lower_shape_func(op, self.target_host)
+        input_states = cfunc.shape_func_param_states
+
+        is_inputs = []
+        input_pos = 0
+        for i, (arg, state) in enumerate(zip(new_args, input_states)):
+            state = int(state)
+            # Pass Shapes
+            if state == 2:
+                for j, subexp in enumerate(from_tuple_type(arg.type_annotation, arg)):
+                    let_in_arg = scope.let("in_arg_{0}".format(input_pos + j), subexp)
+                    sh_of = self.visit(self.shape_of(let_in_arg))
+                    shape_func_ins.append(
+                        scope.let("in_shape_{0}".format(input_pos + j), sh_of))
+                    input_pos += 1
+                is_inputs.append(0)
+            # Pass Inputs
+            elif state == 1:
+                new_arg = self.visit(arg)
+                shape_func_ins.append(
+                    scope.let("in_shape_{0}".format(input_pos), new_arg))
+                input_pos += 1
+                is_inputs.append(1)
+            else:
+                # TODO(@jroesch): handle 3rd case
+                raise Exception("unsupported shape function input state")
+
+        out_shapes = []
+        for i, out in enumerate(cfunc.outputs):
+            tt = ty.TensorType(out.shape, out.dtype)
+            alloc = self.make_static_allocation(scope, tt, i)
+            alloc = scope.let("shape_func_out_{0}".format(i), alloc)
+            out_shapes.append(alloc)
+
+        shape_call = self.shape_func(
+            op,
+            expr.Tuple(shape_func_ins),
+            expr.Tuple(out_shapes), is_inputs)
+
+        scope.let("shape_func", shape_call)
+
+        storages = []
+        for out_shape, out_type in zip(out_shapes, out_types):
+            size = self.compute_storage_in_relay(
+                out_shape, out_type.dtype)
+            alignment = self.compute_alignment(out_type.dtype)
+            sto = scope.let("storage_{i}".format(i=i), self.alloc_storage(
+                size, alignment, self.default_context, out_type.dtype))
+            storages.append(sto)
+
+        outs = []
+        sh_ty_storage = zip(out_shapes, out_types, storages)
+        for i, (out_shape, out_type, storage) in enumerate(sh_ty_storage):
+            alloc = self.alloc_tensor(
+                storage,
+                out_shape,
+                out_type.dtype,
+                out_type.shape)
+            alloc = scope.let("out_{i}".format(i=i), alloc)
+            outs.append(alloc)
+
+        tuple_outs = expr.Tuple(outs)
+        invoke = self.invoke_tvm(op, ins, tuple_outs)
+        scope.let("", invoke)
+        return to_tuple_type(ret_type, tuple_outs.fields)
+
+    def is_dynamic(self, ret_type):
+        is_dynamic = ty.type_has_any(ret_type)
+        # TODO(@jroesch): restore this code, more complex then it seems
+        # for arg in call.args:
+        #     is_dynamic = is_dynamic or arg.checked_type.is_dynamic()
+        return is_dynamic
+
     def visit_call(self, call):
         if is_primitive(call):
             # Because we are in ANF we do not need to visit the arguments.
@@ -131,80 +207,11 @@ def visit_call(self, call):
             ret_type = call.checked_type
             out_types = flatten_tuple_type(ret_type)
 
-            is_dynamic = ty.type_has_any(ret_type)
-            # TODO(@jroesch): restore this code, more complex then it seems
-            # for arg in call.args:
-            #     is_dynamic = is_dynamic or arg.checked_type.is_dynamic()
-
-            if is_dynamic:
-                shape_func_ins = []
-                engine = compile_engine.get()
-                cfunc = engine.lower_shape_func(call.op, self.target_host)
-                input_states = cfunc.shape_func_param_states
-
-                is_inputs = []
-                input_pos = 0
-                for i, (arg, state) in enumerate(zip(new_args, input_states)):
-                    state = int(state)
-                    # Pass Shapes
-                    if state == 2:
-                        for j, subexp in enumerate(from_tuple_type(arg.type_annotation, arg)):
-                            let_in_arg = scope.let("in_arg_{0}".format(input_pos + j), subexp)
-                            sh_of = self.visit(self.shape_of(let_in_arg))
-                            shape_func_ins.append(
-                                scope.let("in_shape_{0}".format(input_pos + j), sh_of))
-                            input_pos += 1
-                        is_inputs.append(0)
-                    # Pass Inputs
-                    elif state == 1:
-                        new_arg = self.visit(arg)
-                        shape_func_ins.append(
-                            scope.let("in_shape_{0}".format(input_pos), new_arg))
-                        input_pos += 1
-                        is_inputs.append(1)
-                    else:
-                        # TODO(@jroesch): handle 3rd case
-                        raise Exception("unsupported shape function input state")
-
-                out_shapes = []
-                for i, out in enumerate(cfunc.outputs):
-                    tt = ty.TensorType(out.shape, out.dtype)
-                    alloc = self.make_static_allocation(scope, tt, i)
-                    alloc = scope.let("shape_func_out_{0}".format(i), alloc)
-                    out_shapes.append(alloc)
-
-                shape_call = self.shape_func(
-                    call.op,
-                    expr.Tuple(shape_func_ins),
-                    expr.Tuple(out_shapes), is_inputs)
-
-                scope.let("shape_func", shape_call)
-
-                storages = []
-                for out_shape, out_type in zip(out_shapes, out_types):
-                    size = self.compute_storage_in_relay(
-                        out_shape, out_type.dtype)
-                    alignment = self.compute_alignment(out_type.dtype)
-                    sto = scope.let("storage_{i}".format(i=i), self.alloc_storage(
-                        size, alignment, self.default_context, out_type.dtype))
-                    storages.append(sto)
-
-                outs = []
-                sh_ty_storage = zip(out_shapes, out_types, storages)
-                for i, (out_shape, out_type, storage) in enumerate(sh_ty_storage):
-                    alloc = self.alloc_tensor(
-                        storage,
-                        out_shape,
-                        out_type.dtype,
-                        out_type.shape)
-                    alloc = scope.let("out_{i}".format(i=i), alloc)
-                    outs.append(alloc)
-
-                tuple_outs = expr.Tuple(outs)
-                invoke = self.invoke_tvm(call.op, ins, tuple_outs)
-                scope.let("", invoke)
-                return to_tuple_type(ret_type, tuple_outs.fields)
+            if self.is_dynamic(ret_type):
+                # Handle dynamic case.
+                return self.dynamic_invoke(scope, op, ins, new_args, out_types, ret_type)
             else:
+                # Handle static case.
                 outs = []
                 for i, out_ty in enumerate(out_types):
                     out = self.make_static_allocation(scope, out_ty, i)

From 967c16f5ce1ed6a995492246a1a84b3dde8476a0 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 13 Apr 2020 17:08:43 -0700
Subject: [PATCH 05/10] Add doc strings

---
 python/tvm/relay/op/memory/memory.py | 41 ++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/python/tvm/relay/op/memory/memory.py b/python/tvm/relay/op/memory/memory.py
index b3488c6d93dd..de4e03a38eb3 100644
--- a/python/tvm/relay/op/memory/memory.py
+++ b/python/tvm/relay/op/memory/memory.py
@@ -101,10 +101,51 @@ def shape_func(func, inputs, outputs, dependent=False):
     return _make.shape_func(func, inputs, outputs, dependent)
 
 def flatten_tuple_type(ty):
+    """Return a sequence of the types contained in the tuple type in order.
+
+    Parameters
+    ----------
+    ty: tvm.Type
+        The type to flatten.
+
+    Returns
+    -------
+    result: List[tvm.Type]
+        The types in their linear order.
+    """
     return _make.FlattenTupleType(ty)
 
 def from_tuple_type(ty, expr):
+    """Convert an expression with the given type into a sequence of expressions.
+       Each expressions maps to a field of the tuple or nested tuples in linear
+       order.
+
+    Parameters
+    ----------
+    ty: tvm.Type
+        The type to unpack.
+    expr: The expression from which to extract each sub-field.
+
+    Returns
+    -------
+    result: List[tvm.relay.Expr]
+        The list of sub-expressions.
+    """
     return _make.FromTupleType(ty, expr)
 
 def to_tuple_type(ty, exprs):
+    """Pack the sequence of expressions into the nested tuple type.
+
+    Parameters
+    ----------
+    ty: tvm.Type
+        The type to pack with.
+
+    exprs: The expressions to pack back into the nested tuple type.
+
+    Returns
+    -------
+    result: List[tvm.relay.Expr]
+    The packed tuple expression.
+    """
     return _make.ToTupleType(ty, exprs)

From 7e324620ed995651eb246af834b70fc61244cf87 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 13 Apr 2020 17:13:49 -0700
Subject: [PATCH 06/10] Fix CPPlint

---
 include/tvm/relay/attrs/memory.h |  7 ++---
 src/relay/op/memory/memory.cc    | 52 +++++++++++++++-----------------
 2 files changed, 28 insertions(+), 31 deletions(-)

diff --git a/include/tvm/relay/attrs/memory.h b/include/tvm/relay/attrs/memory.h
index db4e04eba54f..d232f867a777 100644
--- a/include/tvm/relay/attrs/memory.h
+++ b/include/tvm/relay/attrs/memory.h
@@ -27,6 +27,7 @@
 #include <tvm/ir/attrs.h>
 #include <tvm/relay/expr.h>
 #include <string>
+#include <vector>
 
 namespace tvm {
 namespace relay {
@@ -50,12 +51,10 @@ struct AllocStorageAttrs : public tvm::AttrsNode<AllocStorageAttrs> {
       .set_default(DataType::Float(32, 1));
     TVM_ATTR_FIELD(device_id)
       .describe(
-        "The device id on which to allocate memory."
-      );
+        "The device id on which to allocate memory.");
     TVM_ATTR_FIELD(device_type)
       .describe(
-        "The device type on which to allocate memory."
-      );
+        "The device type on which to allocate memory.");
   }
 };
 
diff --git a/src/relay/op/memory/memory.cc b/src/relay/op/memory/memory.cc
index 54ff3bf29704..3eabbafab8cd 100644
--- a/src/relay/op/memory/memory.cc
+++ b/src/relay/op/memory/memory.cc
@@ -23,13 +23,13 @@
  */
 
 #include <topi/elemwise.h>
+#include <tvm/relay/attrs/memory.h>
 #include <tvm/relay/expr.h>
 #include <tvm/relay/op.h>
 #include <tvm/relay/op_attr_types.h>
-#include <tvm/relay/attrs/memory.h>
 
-#include "../op_common.h"
 #include "../../transforms/infer_layout_util.h"
+#include "../op_common.h"
 #include "../type_relations.h"
 
 namespace tvm {
@@ -91,29 +91,28 @@ RELAY_REGISTER_OP("memory.alloc_storage")
                            });
 
 TVM_REGISTER_GLOBAL("relay.op.memory._make.alloc_tensor")
-    .set_body_typed(
-        [](Expr storage, tvm::relay::Expr shape, DataType dtype, Array<IndexExpr> assert_shape) {
-          auto attrs = make_object<AllocTensorAttrs>();
-          attrs->dtype = dtype;
-          if (assert_shape.defined()) {
-            attrs->assert_shape = assert_shape;
-          } else {
-            attrs->const_shape = Downcast<Constant>(shape);
-          }
-          static const Op& op = Op::Get("memory.alloc_tensor");
-          return Call(op, {storage, shape}, Attrs(attrs), {});
-        });
+    .set_body_typed([](Expr storage, tvm::relay::Expr shape, DataType dtype,
+                       Array<IndexExpr> assert_shape) {
+      auto attrs = make_object<AllocTensorAttrs>();
+      attrs->dtype = dtype;
+      if (assert_shape.defined()) {
+        attrs->assert_shape = assert_shape;
+      } else {
+        attrs->const_shape = Downcast<Constant>(shape);
+      }
+      static const Op& op = Op::Get("memory.alloc_tensor");
+      return Call(op, {storage, shape}, Attrs(attrs), {});
+    });
 
 std::vector<int64_t> FromConstShape(Constant konst) {
   runtime::NDArray shape = konst->data;
   std::vector<int64_t> raw_shape;
   DLTensor tensor = shape.ToDLPack()->dl_tensor;
   CHECK_EQ(tensor.ndim, 1u);
-  CHECK_EQ(tensor.dtype.code, 0U)
-    << "found " << tensor.dtype.code;
+  CHECK_EQ(tensor.dtype.code, 0U) << "found " << tensor.dtype.code;
 
   CHECK(tensor.dtype.bits == 64 || tensor.dtype.bits == 32)
-    << "found " << static_cast<int>(tensor.dtype.bits);
+      << "found " << static_cast<int>(tensor.dtype.bits);
 
   if (tensor.dtype.bits == 32) {
     const int32_t* int_ptr = reinterpret_cast<int32_t*>(tensor.data);
@@ -212,10 +211,9 @@ bool InvokeTVMOPRel(const Array<Type>& types, int num_inputs, const Attrs& attrs
 }
 
 TVM_REGISTER_GLOBAL("relay.op.memory._make.invoke_tvm_op")
-.set_body_typed(
-        [](Expr func, Expr inputs, Expr outputs) {
-          return Call(Op::Get("memory.invoke_tvm_op"), {func, inputs, outputs}, Attrs());
-        });
+    .set_body_typed([](Expr func, Expr inputs, Expr outputs) {
+      return Call(Op::Get("memory.invoke_tvm_op"), {func, inputs, outputs}, Attrs());
+    });
 
 RELAY_REGISTER_OP("memory.invoke_tvm_op")
     .describe(R"code(Invoke an operation compiled by TVM.)code" TVM_ADD_FILELINE)
@@ -260,8 +258,7 @@ RELAY_REGISTER_OP("memory.kill")
                            });
 
 TVM_REGISTER_GLOBAL("relay.op.memory._make.shape_func")
-    .set_body_typed(
-      [](Expr func, Expr inputs, Expr outputs, Array<tvm::Integer> is_input) {
+    .set_body_typed([](Expr func, Expr inputs, Expr outputs, Array<tvm::Integer> is_input) {
       static const Op& op = Op::Get("memory.shape_func");
       auto attrs = make_object<ShapeFuncAttrs>();
       attrs->is_input = is_input;
@@ -304,7 +301,8 @@ std::vector<Expr> FromTupleType(const Type& type, const Expr& expr) {
   return out;
 }
 
-static void ToTupleTypeAux(const Type& type, const std::vector<Expr>& exprs, int* index, std::vector<Expr>* out) {
+static void ToTupleTypeAux(const Type& type, const std::vector<Expr>& exprs, int* index,
+                           std::vector<Expr>* out) {
   if (auto tt = type.as<TensorTypeNode>()) {
     out->push_back(exprs[*index]);
     *index += 1;
@@ -344,9 +342,9 @@ TVM_REGISTER_GLOBAL("relay.op.memory._make.FromTupleType")
 });
 
 TVM_REGISTER_GLOBAL("relay.op.memory._make.ToTupleType")
-.set_body_typed([](Type t, Array<Expr> array) {
-  return ToTupleType(t, std::vector<Expr>(array.begin(), array.end()));
-});
+    .set_body_typed([](Type t, Array<Expr> array) {
+      return ToTupleType(t, std::vector<Expr>(array.begin(), array.end()));
+    });
 
 bool ShapeFuncRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
                   const TypeReporter& reporter) {

From 327c3a818b9895a4d2ea15729173fbf3817b445f Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 13 Apr 2020 17:18:17 -0700
Subject: [PATCH 07/10] Fix PyLint

---
 python/tvm/relay/op/memory/memory.py       | 1 +
 python/tvm/relay/transform/memory_alloc.py | 9 +++++----
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/python/tvm/relay/op/memory/memory.py b/python/tvm/relay/op/memory/memory.py
index de4e03a38eb3..3c01c536ffb9 100644
--- a/python/tvm/relay/op/memory/memory.py
+++ b/python/tvm/relay/op/memory/memory.py
@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=no-else-return,invalid-name,len-as-condition,too-many-nested-blocks
 """Operators for manipulating low-level memory."""
 from __future__ import absolute_import as _abs
 from . import _make
diff --git a/python/tvm/relay/transform/memory_alloc.py b/python/tvm/relay/transform/memory_alloc.py
index db5df4181964..4f9604c02653 100644
--- a/python/tvm/relay/transform/memory_alloc.py
+++ b/python/tvm/relay/transform/memory_alloc.py
@@ -122,10 +122,11 @@ def visit_let(self, let):
 
         return scope.get()
 
-    def dynamic_invoke(self, scope, op, ins, new_args, out_types, ret_type):
+    def dynamic_invoke(self, scope, func, ins, new_args, out_types, ret_type):
+        """Generate the code for invoking a TVM op with a dynamic shape."""
         shape_func_ins = []
         engine = compile_engine.get()
-        cfunc = engine.lower_shape_func(op, self.target_host)
+        cfunc = engine.lower_shape_func(func, self.target_host)
         input_states = cfunc.shape_func_param_states
 
         is_inputs = []
@@ -160,7 +161,7 @@ def dynamic_invoke(self, scope, op, ins, new_args, out_types, ret_type):
             out_shapes.append(alloc)
 
         shape_call = self.shape_func(
-            op,
+            func,
             expr.Tuple(shape_func_ins),
             expr.Tuple(out_shapes), is_inputs)
 
@@ -187,7 +188,7 @@ def dynamic_invoke(self, scope, op, ins, new_args, out_types, ret_type):
             outs.append(alloc)
 
         tuple_outs = expr.Tuple(outs)
-        invoke = self.invoke_tvm(op, ins, tuple_outs)
+        invoke = self.invoke_tvm(func, ins, tuple_outs)
         scope.let("", invoke)
         return to_tuple_type(ret_type, tuple_outs.fields)
 

From 5c7d0ed977e583d25837f31262cbee63812f4960 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 13 Apr 2020 17:24:34 -0700
Subject: [PATCH 08/10] Fix

---
 src/relay/op/memory/memory.cc | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/relay/op/memory/memory.cc b/src/relay/op/memory/memory.cc
index 3eabbafab8cd..0a7142df572f 100644
--- a/src/relay/op/memory/memory.cc
+++ b/src/relay/op/memory/memory.cc
@@ -284,10 +284,10 @@ std::vector<TensorType> FlattenTupleType(const Type& type) {
 }
 
 static void FromTupleTypeAux(const Type& type, const Expr& expr, std::vector<Expr>* out) {
-  if (auto tt = type.as<TensorTypeNode>()) {
+  if (type.as<TensorTypeNode>()) {
     out->push_back(expr);
   } else if (auto tuple_ty = type.as<TupleTypeNode>()) {
-    for (auto i = 0; i < tuple_ty->fields.size(); i++) {
+    for (size_t i = 0; i < tuple_ty->fields.size(); i++) {
       FromTupleTypeAux(tuple_ty->fields[i], TupleGetItem(expr, i), out);
     }
   } else {
@@ -303,12 +303,12 @@ std::vector<Expr> FromTupleType(const Type& type, const Expr& expr) {
 
 static void ToTupleTypeAux(const Type& type, const std::vector<Expr>& exprs, int* index,
                            std::vector<Expr>* out) {
-  if (auto tt = type.as<TensorTypeNode>()) {
+  if (type.as<TensorTypeNode>()) {
     out->push_back(exprs[*index]);
     *index += 1;
   } else if (auto tuple_ty = type.as<TupleTypeNode>()) {
     std::vector<Expr> tuple_out;
-    for (auto i = 0; i < tuple_ty->fields.size(); i++) {
+    for (size_t i = 0; i < tuple_ty->fields.size(); i++) {
       ToTupleTypeAux(tuple_ty->fields[i], exprs, index, &tuple_out);
     }
     out->push_back(Tuple(tuple_out));

From 957aa743cb7e691476ec66394dc1aa68f42463f9 Mon Sep 17 00:00:00 2001
From: Jared Roesch <roeschinc@gmail.com>
Date: Mon, 13 Apr 2020 18:35:59 -0700
Subject: [PATCH 09/10] Apply suggestions from code review

Co-Authored-By: Zhi <5145158+zhiics@users.noreply.github.com>
---
 python/tvm/relay/op/memory/memory.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/python/tvm/relay/op/memory/memory.py b/python/tvm/relay/op/memory/memory.py
index 3c01c536ffb9..509db354b42c 100644
--- a/python/tvm/relay/op/memory/memory.py
+++ b/python/tvm/relay/op/memory/memory.py
@@ -118,14 +118,16 @@ def flatten_tuple_type(ty):
 
 def from_tuple_type(ty, expr):
     """Convert an expression with the given type into a sequence of expressions.
-       Each expressions maps to a field of the tuple or nested tuples in linear
+       Each expression maps to a field of the tuple or nested tuples in linear
        order.
 
     Parameters
     ----------
     ty: tvm.Type
         The type to unpack.
-    expr: The expression from which to extract each sub-field.
+
+    expr: tvm.relay.Expr
+        The expression from which to extract each sub-field.
 
     Returns
     -------
@@ -142,11 +144,12 @@ def to_tuple_type(ty, exprs):
     ty: tvm.Type
         The type to pack with.
 
-    exprs: The expressions to pack back into the nested tuple type.
+    exprs: tvm.relay.Expr
+        The expressions to pack back into the nested tuple type.
 
     Returns
     -------
     result: List[tvm.relay.Expr]
-    The packed tuple expression.
+        The packed tuple expression.
     """
     return _make.ToTupleType(ty, exprs)

From 3da1c0d6833c3be9a47eb1a75d7a92b47c1cdf11 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 13 Apr 2020 18:36:15 -0700
Subject: [PATCH 10/10] Fix typo

---
 python/tvm/relay/transform/memory_alloc.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/tvm/relay/transform/memory_alloc.py b/python/tvm/relay/transform/memory_alloc.py
index 4f9604c02653..611fb1babf55 100644
--- a/python/tvm/relay/transform/memory_alloc.py
+++ b/python/tvm/relay/transform/memory_alloc.py
@@ -210,7 +210,7 @@ def visit_call(self, call):
 
             if self.is_dynamic(ret_type):
                 # Handle dynamic case.
-                return self.dynamic_invoke(scope, op, ins, new_args, out_types, ret_type)
+                return self.dynamic_invoke(scope, call.op, ins, new_args, out_types, ret_type)
             else:
                 # Handle static case.
                 outs = []