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.

Changed files:
test/workers/test_virtual.py, test/torch/tensors/test_pointer.py, syft/frameworks/torch/hook_args.py, test/workers/test_worker.py, syft/workers/base.py

diff --git a/syft/frameworks/torch/hook_args.py b/syft/frameworks/torch/hook_args.py
index c2686dc4f..b84efe08a 100644
--- 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 typed_identity(a):
 
     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
index 2bdca0161..88c873cb3 100644
--- a/syft/workers/base.py
+++ b/syft/workers/base.py
@@ -285,57 +285,39 @@ class BaseWorker(AbstractWorker):
         :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):
         """
diff --git a/test/torch/tensors/test_pointer.py b/test/torch/tensors/test_pointer.py
index 7d205a162..2343da49d 100644
--- 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
index 73f7bbbe7..cb3a16ac4 100644
--- 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
index d58483a8f..adb6626ea 100644
--- 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