Repository: OpenMined/PySyft

Issue:
Get .max() working on PointerTensor; multiple ptrs in response
We can't return tuples of tensors when calling remote executions. This means that functions like .max() fail on PointerTensor objects.


Discussion:
```
import torch 
import syft as sy
hook = sy.TorchHook(torch)
bob = sy.VirtualWorker(hook, id="bob")
x = torch.tensor([1,2,3,4,5])
x_ptr = x.send(bob)
print(x_ptr.max())
```

Currently, `x_ptr.max()` returns  `(Wrapper)>[PointerTensor | me:73509977871 -> bob:73509977871]`
and instead, should return `5` right?

@iamtrask I was wondering, if we should build an interface to directly enable all torch tensor functionalities to pointer tensor? Is that something useful? Do we require all torch functionality for pointer tensor? Why do we specifically need max?
max is one between others (torch.split, etc.), basically all functions that can return more than one tensors with fail if they are run remotely as we explicitely always make the assumption the response is a unique tensor.

Resolving patch:
diff --git a/syft/frameworks/torch/hook_args.py b/syft/frameworks/torch/hook_args.py
--- a/syft/frameworks/torch/hook_args.py
+++ b/syft/frameworks/torch/hook_args.py
@@ -1,12 +1,16 @@
 import torch
+import syft as sy
 from syft.exceptions import RemoteTensorFoundError
 from syft.exceptions import PureTorchTensorFoundError
+from .tensors.interpreters import AbstractTensor
 from .tensors.interpreters import PointerTensor
-from .tensors.decorators import LoggingTensor
 from .tensors.interpreters import TorchTensor
 from .tensors.interpreters import FixedPrecisionTensor
 from .tensors.interpreters import AdditiveSharingTensor
 from .tensors.interpreters import MultiPointerTensor
+from .tensors.decorators import LoggingTensor
+
+from typing import Callable, Union, Tuple
 
 hook_method_args_functions = {}
 hook_method_response_functions = {}
@@ -425,80 +429,89 @@ def build_response_hook(response, rules, wrap_type, wrap_args, return_tuple=Fals
     return lambda x: f(lambdas, x)
 
 
-def zero_fold(*a):
+def zero_fold(*a, **k):
     return tuple()
 
 
-def one_fold(return_tuple):
-    def _one_fold(lambdas, args):
-        return lambdas[0](args[0])
+def one_fold(return_tuple, **kwargs):
+    def _one_fold(lambdas, args, **kwargs):
+        return lambdas[0](args[0], **kwargs)
 
     def tuple_one_fold(lambdas, args):
-        return (lambdas[0](args[0]),)
+        return (lambdas[0](args[0], **kwargs),)
 
     return {False: _one_fold, True: tuple_one_fold}[return_tuple]
 
 
-def two_fold(lambdas, args):
-    return lambdas[0](args[0]), lambdas[1](args[1])
+def two_fold(lambdas, args, **kwargs):
+    return lambdas[0](args[0], **kwargs), lambdas[1](args[1], **kwargs)
 
 
-def three_fold(lambdas, args):
-    return lambdas[0](args[0]), lambdas[1](args[1]), lambdas[2](args[2])
+def three_fold(lambdas, args, **kwargs):
+    return (
+        lambdas[0](args[0], **kwargs),
+        lambdas[1](args[1], **kwargs),
+        lambdas[2](args[2], **kwargs),
+    )
 
 
-def four_fold(lambdas, args):
-    return (lambdas[0](args[0]), lambdas[1](args[1]), lambdas[2](args[2]), lambdas[3](args[3]))
+def four_fold(lambdas, args, **kwargs):
+    return (
+        lambdas[0](args[0], **kwargs),
+        lambdas[1](args[1], **kwargs),
+        lambdas[2](args[2], **kwargs),
+        lambdas[3](args[3], **kwargs),
+    )
 
 
-def five_fold(lambdas, args):
+def five_fold(lambdas, args, **kwargs):
     return (
-        lambdas[0](args[0]),
-        lambdas[1](args[1]),
-        lambdas[2](args[2]),
-        lambdas[3](args[3]),
-        lambdas[4](args[4]),
+        lambdas[0](args[0], **kwargs),
+        lambdas[1](args[1], **kwargs),
+        lambdas[2](args[2], **kwargs),
+        lambdas[3](args[3], **kwargs),
+        lambdas[4](args[4], **kwargs),
     )
 
 
-def six_fold(lambdas, args):
+def six_fold(lambdas, args, **kwargs):
     return (
-        lambdas[0](args[0]),
-        lambdas[1](args[1]),
-        lambdas[2](args[2]),
-        lambdas[3](args[3]),
-        lambdas[4](args[4]),
-        lambdas[5](args[5]),
+        lambdas[0](args[0], **kwargs),
+        lambdas[1](args[1], **kwargs),
+        lambdas[2](args[2], **kwargs),
+        lambdas[3](args[3], **kwargs),
+        lambdas[4](args[4], **kwargs),
+        lambdas[5](args[5], **kwargs),
     )
 
 
-def seven_fold(lambdas, args):
+def seven_fold(lambdas, args, **kwargs):
     return (
-        lambdas[0](args[0]),
-        lambdas[1](args[1]),
-        lambdas[2](args[2]),
-        lambdas[3](args[3]),
-        lambdas[4](args[4]),
-        lambdas[5](args[5]),
-        lambdas[6](args[6]),
+        lambdas[0](args[0], **kwargs),
+        lambdas[1](args[1], **kwargs),
+        lambdas[2](args[2], **kwargs),
+        lambdas[3](args[3], **kwargs),
+        lambdas[4](args[4], **kwargs),
+        lambdas[5](args[5], **kwargs),
+        lambdas[6](args[6], **kwargs),
     )
 
 
-def eight_fold(lambdas, args):
+def eight_fold(lambdas, args, **kwargs):
     return (
-        lambdas[0](args[0]),
-        lambdas[1](args[1]),
-        lambdas[2](args[2]),
-        lambdas[3](args[3]),
-        lambdas[4](args[4]),
-        lambdas[5](args[5]),
-        lambdas[6](args[6]),
-        lambdas[7](args[7]),
+        lambdas[0](args[0], **kwargs),
+        lambdas[1](args[1], **kwargs),
+        lambdas[2](args[2], **kwargs),
+        lambdas[3](args[3], **kwargs),
+        lambdas[4](args[4], **kwargs),
+        lambdas[5](args[5], **kwargs),
+        lambdas[6](args[6], **kwargs),
+        lambdas[7](args[7], **kwargs),
     )
 
 
-def many_fold(lambdas, args):
-    return tuple([lambdas[i](args[i]) for i in range(len(lambdas))])
+def many_fold(lambdas, args, **kwargs):
+    return tuple([lambdas[i](args[i], **kwargs) for i in range(len(lambdas))])
 
 
 # Add the possibility to make a type check in the identity function applied
@@ -531,3 +544,150 @@ def number_identity(i):
 
     else:
         return lambda i: i
+
+
+# -- Fast way to register responses and transform tensors in pointers
+
+register_response_functions = {}
+
+
+def register_response(attr: str, response: object, owner: sy.workers.AbstractWorker) -> object:
+    """
+    When a remote worker execute a command sent by someone else, the response is
+    inspected: all tensors are stored by this worker and a Pointer tensor is
+    made for each of them.
+
+    To make this efficient, we cache which elements of the response (which can be more
+    complicated with nested tuples for example) in the dict register_response_functions
+
+    However, sometimes a function  (an attr) has multiple different response signatures.
+    This invalidates the cache, so we need to have a try/except which refreshes the
+    cache if the signature triggers an error.
+
+    Args:
+        attr (str): the name of the function being called
+        response (object): the response of this function
+        owner (BaseWorker): the worker which registers the tensors
+    """
+
+    # TODO: Why do we need to cast it in a tuple? this is a (small) time waste
+    response_is_tuple = isinstance(response, tuple)
+
+    # Add an artificial tuple
+    if not response_is_tuple:
+        response = (response, 1)
+
+    attr_id = "{}".format(attr)
+
+    try:
+        # Load the utility function to register the response and transform tensors with pointers
+        register_response_function = register_response_functions[attr_id]
+        # Try running it
+        new_response = register_response_function(response, owner=owner)
+
+    except (IndexError, KeyError, AssertionError):  # Update the function in cas of an error
+        register_response_function = build_register_response_function(response)
+        # Store this utility function in the registry
+        register_response_functions[attr_id] = register_response_function
+        # Run it
+        new_response = register_response_function(response, owner=owner)
+
+    # Remove the artificial tuple
+    if not response_is_tuple:
+        new_response, _ = new_response
+
+    return new_response
+
+
+def build_register_response_function(response: object) -> Callable:
+    """
+    Build the function that registers the response and replaces tensors with pointers.
+
+    Example:
+        (1, tensor([1, 2]) is the response
+        f is the register_response_function
+        then f(p) = (1, (Wrapper)>Pointer)
+    """
+    # Inspect the call to find tensor arguments and return a rule whose
+    # structure is the same as the response object, with 1 where there was
+    # (torch or syft) tensors and 0 when not (ex: number, str, ...)
+    rule = build_rule(response)
+    # Build a function with this rule to efficiently replace syft tensors
+    # (but not pointer) with their child in the args objects
+    response_hook_function = build_register_response(response, rule)
+    return response_hook_function
+
+
+def register_transform_tensor(
+    tensor: Union[torch.Tensor, AbstractTensor], owner: sy.workers.AbstractWorker = None
+) -> PointerTensor:
+    """
+    Register a tensor and create a pointer that references it
+
+    Args:
+        tensor: the tensor
+        owner: the owner make the registration
+    Returns:
+        the pointer
+    """
+    assert owner is not None
+    # FIXME: should be added automatically
+    tensor.owner = owner
+
+    owner.register_obj(tensor)
+
+    pointer = tensor.create_pointer(
+        location=owner,
+        id_at_location=tensor.id,
+        register=True,
+        owner=owner,
+        ptr_id=tensor.id,
+        garbage_collect_data=False,
+    )
+    return pointer
+
+
+def build_register_response(response: object, rules: Tuple, return_tuple: bool = False) -> Callable:
+    """
+    Build a function given some rules to efficiently replace in the response object
+    torch tensors with a pointer after they are registered, and do nothing for other
+    types of object including , str, numbers, bool, etc.
+
+    Args:
+        response: the response
+        rules: the rule specifying where the tensors are
+        return_tuple: force to return a tuple even with a single element
+    Returns:
+        The function to apply on generic responses
+    """
+
+    # get the transformation lambda for each args
+    lambdas = [
+        (lambda i, **kwargs: i)  # return the same object
+        if not r  # if the rule is a number == 0.
+        else build_register_response(a, r, True)  # If not, call recursively build_response_hook
+        if isinstance(r, (list, tuple))  # if the rule is a list or tuple.
+        # Last if not, rule is probably == 1 so use type to return the right transformation.
+        else lambda i, **kwargs: register_transform_tensor(i, **kwargs)
+        for a, r in zip(response, rules)  # And do this for all the responses / rules provided
+    ]
+
+    # Instead of iterating which is slow, we use trick to efficiently
+    # apply each lambda to each arg
+    folds = {
+        0: zero_fold,
+        1: one_fold(return_tuple),
+        2: two_fold,
+        3: three_fold,
+        4: four_fold,
+        5: five_fold,
+        6: six_fold,
+        7: seven_fold,
+        8: eight_fold,
+    }
+    try:
+        f = folds[len(lambdas)]
+    except KeyError:
+        f = many_fold
+
+    return lambda x, **kwargs: f(lambdas, x, **kwargs)
diff --git a/syft/workers/base.py b/syft/workers/base.py
--- a/syft/workers/base.py
+++ b/syft/workers/base.py
@@ -285,57 +285,39 @@ def execute_command(self, message):
         :return: a pointer to the result
         """
 
-        command, _self, args, kwargs = message
+        command_name, _self, args, kwargs = message
 
         # TODO add kwargs
         kwargs = {}
-        command = command.decode("utf-8")
+        command_name = command_name.decode("utf-8")
         # Handle methods
         if _self is not None:
-            if sy.torch.is_inplace_method(command):
-                getattr(_self, command)(*args, **kwargs)
+            if sy.torch.is_inplace_method(command_name):
+                getattr(_self, command_name)(*args, **kwargs)
                 return
             else:
-                tensor = getattr(_self, command)(*args, **kwargs)
+                response = getattr(_self, command_name)(*args, **kwargs)
         # Handle functions
         else:
             # At this point, the command is ALWAYS a path to a
             # function (i.e., torch.nn.functional.relu). Thus,
             # we need to fetch this function and run it.
 
-            sy.torch.command_guard(command, "torch_modules")
+            sy.torch.command_guard(command_name, "torch_modules")
 
-            paths = command.split(".")
+            paths = command_name.split(".")
             command = self
             for path in paths:
                 command = getattr(command, path)
 
-            tensor = command(*args, **kwargs)
+            response = command(*args, **kwargs)
 
         # some functions don't return anything (such as .backward())
         # so we need to check for that here.
-        if tensor is not None:
-
-            # FIXME: should be added automatically
-            tensor.owner = self
-
-            # TODO: Handle when the response is not simply a tensor
-            # don't re-register tensors if the operation was inline
-            # not only would this be inefficient, but it can cause
-            # serious issues later on
-            # if(_self is not None):
-            #     if(tensor.id != _self.id):
-            self.register_obj(tensor)
-
-            pointer = tensor.create_pointer(
-                location=self,
-                id_at_location=tensor.id,
-                register=True,
-                owner=self,
-                ptr_id=tensor.id,
-                garbage_collect_data=False,
-            )
-            return pointer
+        if response is not None:
+            # Register response et create pointers for tensor elements
+            response = sy.frameworks.torch.hook_args.register_response(command_name, response, self)
+            return response
 
     def send_command(self, recipient, message):
         """


Test patch:
diff --git a/test/torch/tensors/test_pointer.py b/test/torch/tensors/test_pointer.py
--- a/test/torch/tensors/test_pointer.py
+++ b/test/torch/tensors/test_pointer.py
@@ -259,3 +259,23 @@ def test_remote_to_cpu_device(workers):
 
     x = th.tensor([1, 2, 3, 4, 5]).send(bob)
     x.to(device)
+
+
+def test_remote_function_with_multi_ouput(workers):
+    """
+    Functions like .split return several tensors, registration and response
+    must be made carefully in this case
+    """
+    bob = workers["bob"]
+
+    tensor = torch.tensor([1, 2, 3, 4.0])
+    ptr = tensor.send(bob)
+    r_ptr = torch.split(ptr, 2)
+    assert (r_ptr[0].get() == torch.tensor([1, 2.0])).all()
+
+    tensor = torch.tensor([1, 2, 3, 4.0])
+    ptr = tensor.send(bob)
+    max_value, argmax_idx = torch.max(ptr, 0)
+
+    assert max_value.get().item() == 4.0
+    assert argmax_idx.get().item() == 3
diff --git a/test/workers/test_virtual.py b/test/workers/test_virtual.py
--- a/test/workers/test_virtual.py
+++ b/test/workers/test_virtual.py
@@ -1,5 +1,6 @@
-import syft as sy
+import random
 
+import syft as sy
 from syft.workers.virtual import VirtualWorker
 from syft.codes import MSGTYPE
 from syft import serde
@@ -19,7 +20,8 @@ def test_send_msg():
     me = sy.torch.hook.local_worker
 
     # create a new worker (to send the object to)
-    bob = VirtualWorker(sy.torch.hook)
+    worker_id = int(10e10 * random.random())
+    bob = VirtualWorker(sy.torch.hook, id=f"bob{worker_id}")
 
     # initialize the object and save it's id
     obj = torch.Tensor([100, 100])
@@ -40,7 +42,8 @@ def test_send_msg_using_tensor_api():
     """
 
     # create worker to send object to
-    bob = VirtualWorker(sy.torch.hook)
+    worker_id = int(10e10 * random.random())
+    bob = VirtualWorker(sy.torch.hook, id=f"bob{worker_id}")
 
     # create a tensor to send (default on local_worker)
     obj = torch.Tensor([100, 100])
@@ -66,7 +69,8 @@ def test_recv_msg():
     # TEST 1: send tensor to alice
 
     # create a worker to send data to
-    alice = VirtualWorker(sy.torch.hook)
+    worker_id = int(10e10 * random.random())
+    alice = VirtualWorker(sy.torch.hook, id=f"alice{worker_id}")
 
     # create object to send
     obj = torch.Tensor([100, 100])
@@ -113,8 +117,10 @@ def tests_worker_convenience_methods():
     """
 
     me = sy.torch.hook.local_worker
-    bob = VirtualWorker(sy.torch.hook)
-    alice = VirtualWorker(sy.torch.hook)
+    worker_id = int(10e10 * random.random())
+    bob = VirtualWorker(sy.torch.hook, id=f"bob{worker_id}")
+    worker_id = int(10e10 * random.random())
+    alice = VirtualWorker(sy.torch.hook, id=f"alice{worker_id}")
     obj = torch.Tensor([100, 100])
 
     # Send data to alice
@@ -142,7 +148,8 @@ def tests_worker_convenience_methods():
 
 
 def test_search():
-    bob = VirtualWorker(sy.torch.hook)
+    worker_id = int(10e10 * random.random())
+    bob = VirtualWorker(sy.torch.hook, id=f"bob{worker_id}")
 
     x = (
         torch.tensor([1, 2, 3, 4, 5])
diff --git a/test/workers/test_worker.py b/test/workers/test_worker.py
--- a/test/workers/test_worker.py
+++ b/test/workers/test_worker.py
@@ -1,5 +1,6 @@
 import pytest
 
+import random
 import torch
 import syft as sy
 from syft.exceptions import WorkerNotFoundException
@@ -11,15 +12,20 @@ def test___init__():
 
     tensor = torch.tensor([1, 2, 3, 4])
 
-    alice = VirtualWorker(hook, id="alice")
-    bob = VirtualWorker(hook, id="bob")
-    charlie = VirtualWorker(hook, id="charlie")
-    dawson = VirtualWorker(hook, id="dawson", data=[tensor])
+    worker_id = int(10e10 * random.random())
+    alice_id = f"alice{worker_id}"
+    alice = VirtualWorker(hook, id=alice_id)
+    worker_id = int(10e10 * random.random())
+    bob = VirtualWorker(hook, id=f"bob{worker_id}")
+    worker_id = int(10e10 * random.random())
+    charlie = VirtualWorker(hook, id=f"charlie{worker_id}")
+    worker_id = int(10e10 * random.random())
+    dawson = VirtualWorker(hook, id=f"dawson{worker_id}", data=[tensor])
 
     # Ensure adding data on signup functionality works as expected
     assert tensor.owner == dawson
 
-    assert bob.get_worker("alice").id == alice.id
+    assert bob.get_worker(alice_id).id == alice.id
     assert bob.get_worker(alice).id == alice.id
     assert bob.get_worker(charlie).id == charlie.id