Improve Plans flexibility

syft/execution/plan.py

@@ -232,9 +232,15 @@ def add_placeholder(self, tensor, node_type=None):
                         f"Please use instead torch.tensor(..., dtype=torch.int32) for example."
                     )
                 placeholder.tags.add(f"#input-{self._tmp_args_ids.index(tensor.id)}")
+                if tensor.id in self._tmp_result_ids:
+                    placeholder.tags.add(f"#output-{self._tmp_result_ids.index(tensor.id)}")
+
             elif node_type == "output":
                 if tensor.id in self._tmp_result_ids:
                     placeholder.tags.add(f"#output-{self._tmp_result_ids.index(tensor.id)}")
+
+                if tensor.id in self._tmp_args_ids:
+                    placeholder.tags.add(f"#input-{self._tmp_result_ids.index(tensor.id)}")
             else:
                 raise ValueError("node_type should be 'input' or 'output'.")
 
@@ -318,6 +324,9 @@ def build(self, *args):
         results = (results,) if not isinstance(results, tuple) else results
         self._tmp_result_ids = [t.id for t in results if isinstance(t, FrameworkTensor)]
 
+        for arg in args:
+            self.replace_with_placeholders(arg, node_type="input")
+
         for log in sy.hook.trace.logs:
             command, response = log
             command_placeholders = self.replace_with_placeholders(command, node_type="input")

test/execution/test_plan.py

@@ -1096,3 +1096,33 @@ def foo(x, state):
     y = torchscript_plan(t)
 
     assert (y == th.tensor([3.0, 5])).all()
+
+
+def test_plan_input_usage(hook):
+    x11 = th.tensor([-1, 2.0]).tag("input_data")
+    x12 = th.tensor([1, -2.0]).tag("input_data2")
+
+    device_1 = sy.VirtualWorker(hook, id="test_dev_1", data=(x11, x12))
+
+    @sy.func2plan()
+    def plan_test_1(x, y):
+        return x
+
+    @sy.func2plan()
+    def plan_test_2(x, y):
+        return y
+
+    pointer_to_data_1 = device_1.search("input_data")[0]
+    pointer_to_data_2 = device_1.search("input_data2")[0]
+
+    plan_test_1.build(th.tensor([1.0, -2.0]), th.tensor([1, 2]))
+    pointer_plan = plan_test_1.send(device_1)
+    pointer_to_result = pointer_plan(pointer_to_data_1, pointer_to_data_2)
+    result = pointer_to_result.get()
+    assert (result == x11).all()
+
+    plan_test_2.build(th.tensor([1.0, -2.0]), th.tensor([1, 2]))
+    pointer_plan = plan_test_2.send(device_1)
+    pointer_to_result = pointer_plan(pointer_to_data_1, pointer_to_data_2)
+    result = pointer_to_result.get()
+    assert (result == x12).all