add set_one_input method to python API of VirtualMachine

Author: Valery Chernov

python/tvm/runtime/vm.py

@@ -380,6 +380,7 @@ def __init__(self, exe, device, memory_cfg=None):
         self._get_num_outputs = self.module["get_num_outputs"]
         self._get_input_index = self.module["get_input_index"]
         self._set_input = self.module["set_input"]
+        self._set_one_input = self.module["set_one_input"]
         self._setup_device(device, memory_cfg)
 
     def _setup_device(self, dev, memory_cfg):
@@ -450,6 +451,24 @@ def set_input(self, func_name, *args, **kwargs):
         cargs = convert(args)
         self._set_input(func_name, *cargs)
 
+    def set_one_input(self, func_name, **kwargs):
+        """Set the one input tensor with tag to a function.
+
+        Parameters
+        ----------
+        func_name : str
+            The name of the function.
+
+        kwargs: dict of str or int to tvm.runtime.NDArray or np.ndarray
+            Named arguments to the function.
+        """
+        assert len(kwargs) == 1
+        tag = kwargs.keys()[0]
+        if isinstance(tag, str):
+            func_params = self._exec.get_function_params(func_name)
+            assert tag in func_params
+        self._set_one_input(func_name, tag, kwargs[tag])
+
     def invoke(self, func_name, *args, **kwargs):
         """Invoke a function.