Add mod supoort in relay.build

python/tvm/relay/build_module.py

@@ -143,8 +143,8 @@ def build(func, target=None, target_host=None, params=None):
 
     Parameters
     ----------
-    func: relay.Function
-        The function to build.
+    func: relay.Function or relay.Module
+        The function or module to build.
 
     target : str, :any:`tvm.target.Target`, or dict of str(i.e. device/context
     name) to str/tvm.target.Target, optional
@@ -192,6 +192,9 @@ def build(func, target=None, target_host=None, params=None):
 
     with tophub_context:
         bld_mod = BuildModule()
+        if "tvm.relay.module.Module" in str(type(func)):
+            mod = func
+            func = mod[mod.entry_func]
         graph_json, mod, params = bld_mod.build(func, target, target_host,
                                                 params)
     return graph_json, mod, params

tests/python/frontend/tflite/test_forward.py

@@ -68,9 +68,7 @@ def run_tvm_graph(tflite_model_buf, input_data, input_node, num_output=1, target
                                              shape_dict=shape_dict,
                                              dtype_dict=dtype_dict)
     with relay.build_config(opt_level=3):
-        graph, lib, params = relay.build(mod[mod.entry_func],
-                                         target,
-                                         params=params)
+        graph, lib, params = relay.build(mod, target, params=params)
 
     ctx = tvm.context(target, 0)
     from tvm.contrib import graph_runtime

tutorials/frontend/from_tflite.py

@@ -145,7 +145,7 @@ def extract(path):
 # target x86 CPU
 target = "llvm"
 with relay.build_config(opt_level=3):
-    graph, lib, params = relay.build(mod[mod.entry_func], target, params=params)
+    graph, lib, params = relay.build(mod, target, params=params)
 
 ######################################################################
 # Execute on TVM