Repository: OpenMined/PySyft

Issue:
Make it possible for `Protocols` to generate tensors on specific `Roles`/`Workers`
**Is your feature request related to a problem? Please describe.**
When tracing and executing `Protocols`, it's not yet possible to locally generate new tensors (e.g. with `torch.rand()`) and have them registered on the correct `Role`/`Workers`.

**Describe the solution you'd like**
Give each `Role` a `torch` attribute that contains a `FrameworkWrapper` around the actual `torch` module for tracing into that `Role`.

**Describe alternatives you've considered**
* I suppose we could instead add an additional `kwarg` for which `Role` to trace into onto all `torch` methods and strip it out in the `FrameworkWrapper` before forwarding to `torch`. 🤔 

**Additional context**
This may depend on #3490 (not sure.)



Discussion:

Resolving patch:
diff --git a/syft/execution/role.py b/syft/execution/role.py
--- a/syft/execution/role.py
+++ b/syft/execution/role.py
@@ -11,6 +11,7 @@
 from syft.execution.placeholder import PlaceHolder
 from syft.execution.placeholder_id import PlaceholderId
 from syft.execution.state import State
+from syft.execution.tracing import FrameworkWrapper
 from syft.generic.frameworks.types import FrameworkTensor
 from syft.serde.syft_serializable import SyftSerializable
 from syft.workers.abstract import AbstractWorker
@@ -34,7 +35,8 @@ def __init__(
         output_placeholder_ids: Tuple[int, str] = None,
     ):
         self.id = id or sy.ID_PROVIDER.pop()
-        self.worker = worker
+        self.worker = worker or sy.local_worker
+
         self.actions = actions or []
 
         # All placeholders
@@ -47,6 +49,10 @@ def __init__(
         self.state = state or State()
         self.tracing = False
 
+        for name, package in framework_packages.items():
+            tracing_wrapper = FrameworkWrapper(package=package, role=self, owner=self.worker)
+            setattr(self, name, tracing_wrapper)
+
     def input_placeholders(self):
         return [self.placeholders[id_] for id_ in self.input_placeholder_ids]
 


Test patch:
diff --git a/test/execution/test_protocol.py b/test/execution/test_protocol.py
--- a/test/execution/test_protocol.py
+++ b/test/execution/test_protocol.py
@@ -8,7 +8,7 @@
 def test_func2protocol_creates_roles():
     @sy.func2protocol(roles=["alice", "bob"])
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         return tensor
 
@@ -18,10 +18,25 @@ def protocol(alice, bob):
     assert isinstance(protocol.roles["bob"], Role)
 
 
+def test_framework_methods_traced_by_role():
+    @sy.func2protocol(roles=["alice", "bob"])
+    def protocol(alice, bob):
+        tensor1 = alice.torch.rand([4, 4])
+        tensor2 = bob.torch.rand([4, 4])
+
+        return tensor1, tensor2
+
+    assert protocol.is_built
+
+    for role in protocol.roles.values():
+        assert len(role.actions) == 1
+        assert "torch.rand" in [action.name for action in role.actions]
+
+
 def test_trace_communication_actions_send():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),)})
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         tensor.send(bob.worker)
         return tensor
@@ -29,14 +44,14 @@ def protocol(alice, bob):
     traced_actions = protocol.roles["alice"].actions
 
     assert protocol.is_built
-    assert len(traced_actions) == 1
+    assert len(traced_actions) == 2
     assert "send" in [action.name for action in traced_actions]
 
 
 def test_trace_communication_actions_get():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),)})
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         ptr = tensor.send(bob.worker)
         res = ptr.get()
@@ -45,14 +60,14 @@ def protocol(alice, bob):
     traced_actions = protocol.roles["alice"].actions
 
     assert protocol.is_built
-    assert len(traced_actions) == 2
+    assert len(traced_actions) == 3
     assert "get" in [action.name for action in traced_actions]
 
 
 def test_trace_communication_actions_ptr_send():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),)})
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         ptr = tensor.send(bob.worker)
         res = ptr.send(alice.worker)
@@ -61,14 +76,14 @@ def protocol(alice, bob):
     traced_actions = protocol.roles["alice"].actions
 
     assert protocol.is_built
-    assert len(traced_actions) == 2
+    assert len(traced_actions) == 3
     assert "send" in [action.name for action in traced_actions]
 
 
 def test_trace_communication_actions_move():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),)})
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         ptr = tensor.send(bob.worker)
         res = ptr.move(alice.worker)
@@ -77,14 +92,14 @@ def protocol(alice, bob):
     traced_actions = protocol.roles["alice"].actions
 
     assert protocol.is_built
-    assert len(traced_actions) == 2
+    assert len(traced_actions) == 3
     assert "move" in [action.name for action in traced_actions]
 
 
 def test_trace_communication_actions_share():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),)})
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         ptr = tensor.send(bob.worker)
         ptr = ptr.fix_prec()
@@ -94,14 +109,14 @@ def protocol(alice, bob):
     traced_actions = protocol.roles["alice"].actions
 
     assert protocol.is_built
-    assert len(traced_actions) == 3
+    assert len(traced_actions) == 4
     assert "share" in [action.name for action in traced_actions]
 
 
 def test_trace_communication_actions_share_():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),)})
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         ptr = tensor.send(bob.worker)
         ptr = ptr.fix_prec()
@@ -111,14 +126,14 @@ def protocol(alice, bob):
     traced_actions = protocol.roles["alice"].actions
 
     assert protocol.is_built
-    assert len(traced_actions) == 3
+    assert len(traced_actions) == 4
     assert "share_" in [action.name for action in traced_actions]
 
 
 def test_trace_communication_actions_remote_send():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),)})
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         ptr = tensor.send(bob.worker)
         res = ptr.remote_send(alice.worker)
@@ -127,14 +142,14 @@ def protocol(alice, bob):
     traced_actions = protocol.roles["alice"].actions
 
     assert protocol.is_built
-    assert len(traced_actions) == 2
+    assert len(traced_actions) == 3
     assert "remote_send" in [action.name for action in traced_actions]
 
 
 def test_trace_communication_actions_mid_get():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),)})
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         ptr = tensor.send(bob.worker)
         res = ptr.mid_get()
@@ -143,14 +158,14 @@ def protocol(alice, bob):
     traced_actions = protocol.roles["alice"].actions
 
     assert protocol.is_built
-    assert len(traced_actions) == 2
+    assert len(traced_actions) == 3
     assert "mid_get" in [action.name for action in traced_actions]
 
 
 def test_trace_communication_actions_remote_get():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),)})
     def protocol(alice, bob):
-        tensor = alice.fetch(th.tensor([1]))
+        tensor = alice.torch.tensor([1])
 
         ptr = tensor.send(bob.worker).send(alice.worker)
         res = ptr.remote_get()
@@ -159,15 +174,15 @@ def protocol(alice, bob):
     traced_actions = protocol.roles["alice"].actions
 
     assert protocol.is_built
-    assert len(traced_actions) == 3
+    assert len(traced_actions) == 4
     assert "remote_get" in [action.name for action in traced_actions]
 
 
 def test_create_roles_from_decorator():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),), "bob": ((1,),)})
     def protocol(alice, bob):
-        tensor1 = alice.fetch(th.tensor([1]))
-        tensor2 = bob.fetch(th.tensor([2]))
+        tensor1 = alice.torch.tensor([1])
+        tensor2 = bob.torch.tensor([2])
 
         t1plus = tensor1 + 1
         t2plus = tensor2 + 1
@@ -182,8 +197,8 @@ def protocol(alice, bob):
 def test_multi_role_tracing():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),), "bob": ((1,),)})
     def protocol(alice, bob):
-        tensor1 = alice.fetch(th.tensor([1]))
-        tensor2 = bob.fetch(th.tensor([2]))
+        tensor1 = alice.torch.tensor([1])
+        tensor2 = bob.torch.tensor([2])
 
         t1plus = tensor1 + 1
         t2plus = tensor2 + 1
@@ -195,16 +210,16 @@ def protocol(alice, bob):
     assert protocol.is_built
 
     assert len(protocol.roles) == 2
-    assert len(protocol.roles["alice"].actions) == 1
-    assert len(protocol.roles["bob"].actions) == 1
+    assert len(protocol.roles["alice"].actions) == 2
+    assert len(protocol.roles["bob"].actions) == 2
 
 
 def test_multi_role_execution():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,), (1,)), "bob": ((1,),)})
     def protocol(alice, bob):
-        tensor1 = alice.fetch(th.tensor([1]))
-        tensor2 = bob.fetch(th.tensor([2]))
-        tensor3 = alice.fetch(th.tensor([3]))
+        tensor1 = alice.torch.tensor([1])
+        tensor2 = bob.torch.tensor([2])
+        tensor3 = alice.torch.tensor([3])
 
         res1 = tensor2
         res2 = tensor1 + tensor3
@@ -225,8 +240,8 @@ def protocol(alice, bob):
 def test_copy():
     @sy.func2protocol(roles=["alice", "bob"], args_shape={"alice": ((1,),), "bob": ((1,),)})
     def protocol(alice, bob):
-        tensor1 = alice.fetch(th.tensor([1]))
-        tensor2 = bob.fetch(th.tensor([2]))
+        tensor1 = alice.torch.tensor([1])
+        tensor2 = bob.torch.tensor([2])
 
         t1plus = tensor1 + 1
         t2plus = tensor2 + 1