Repository: OpenMined/PySyft
Pull request: Plans support for typechecking, lists and dicts using serialized input.
# Description

This PR fixes #3184. 

New features:
- Plans now have a new attribute: serialized_input, that will serialized the input at build time and compare it with the serialized input at call time.
- Plans now have support for type-checking at runtime based on the type provided at build time.
- Plans now support usage of lists, tuples and dicts as inputs.
- Plans now have a number of arguments check, build time and call time should have the same number of arguments.

# Type of change:
- New feature

# How Has This Been Tested?

Wrote a few tests that verify that:
- type-checking works, it should raise a RunTimeWarning if the types differ.
- argument number check works, build time and call time should not differ. If they differ, a TypeError should occur. I have chosen TypeError to be consistent with python functions.
- dicts from build time have the same keys as dicts from call time (nested as well). If they differ, a KeyError should occur.
- lists, dicts and tuples of tensors work as input (and any nested combination).

# TODO:
~~Adapt the changes to the new Role class.~~
~~Add syft-proto support.~~
~~Add protobuffer support.~~
~~Remove the usage of pickle (only a PoC). ~~
~~Check interaction with torchscript. ~~

Problem statement:
Improve Plans flexibility
**Is your feature request related to a problem? Please describe.**
Plans currently don't work exactly as normal function in the following scenarios:
- if some inputs that are not used in the computations
- if some inputs that are not used and are given in the output
- if input/ouput contain nested objects like list or tuple

**Describe the solution you'd like**
We want to support this by default

Hints from issue discussion:
I will have a look on this.

Commit message:  Adding docstrings.

Patch to review:
diff --git a/pip-dep/requirements.txt b/pip-dep/requirements.txt
index 60a812ac690..883ed5268d9 100644
--- a/pip-dep/requirements.txt
+++ b/pip-dep/requirements.txt
@@ -7,7 +7,7 @@ phe~=1.4.0
 Pillow~=6.2.2
 requests~=2.22.0
 scipy~=1.4.1
-syft-proto~=0.4.1
+git+git://github.com/tudorcebere/syft-proto@8e60132#egg=syft-proto
 tblib~=1.6.0
 torchvision~=0.5.0
 torch~=1.4.0
diff --git a/syft/execution/plan.py b/syft/execution/plan.py
index 82de40aeab1..c4fe681862e 100644
--- a/syft/execution/plan.py
+++ b/syft/execution/plan.py
@@ -14,6 +14,7 @@
 from syft.execution.role import Role
 from syft.execution.state import State
 from syft.execution.tracing import trace
+from syft.execution.type_wrapper import NestedTypeWrapper
 from syft.execution.translation.abstract import AbstractPlanTranslator
 from syft.execution.translation.default import PlanTranslatorDefault
 from syft.execution.translation.torchscript import PlanTranslatorTorchscript
@@ -114,6 +115,7 @@ def __init__(
         id: Union[str, int] = None,
         owner: "sy.workers.BaseWorker" = None,
         tags: List[str] = None,
+        input_types: list = None,
         description: str = None,
     ):
         AbstractObject.__init__(self, id, owner, tags, description, child=None)
@@ -132,6 +134,7 @@ def __init__(
         self.state_attributes = {}
         self.is_built = is_built
         self.torchscript = None
+        self.input_types = input_types
         self.tracing = False
 
         # The plan has not been sent so it has no reference to remote locations
@@ -179,32 +182,48 @@ def build(self, *args, trace_autograd=False):
             args: Input arguments to run the plan
         """
 
+        def build_nested_arg(arg, leaf_function):
+            if isinstance(arg, list):
+                return [build_nested_arg(obj, leaf_function) for obj in arg]
+            elif isinstance(arg, tuple):
+                return tuple([build_nested_arg(obj, leaf_function) for obj in arg])
+            elif isinstance(arg, dict):
+                return {k: build_nested_arg(v, leaf_function) for k, v in arg.items()}
+            else:
+                return leaf_function(arg)
+
         # Enable tracing
         self.toggle_tracing(True)
         self.is_building = True
 
+        # typecheck
+        self.input_types = NestedTypeWrapper(args)
+
+        # Run once to build the plan
         if trace_autograd:
             # Wrap arguments that require gradients with AutogradTensor,
             # to be able to trace autograd operations
-            args = tuple(
-                AutogradTensor().on(arg, wrap=False)
-                if isinstance(arg, FrameworkTensor) and arg.requires_grad
-                else arg
-                for arg in args
+            args = build_nested_arg(
+                args,
+                lambda x: AutogradTensor().on(x, wrap=False)
+                if isinstance(x, FrameworkTensor) and x.requires_grad
+                else x,
             )
             # Add Placeholder after AutogradTensor in the chain
             # so that all operations that happen inside AutogradTensor are recorded by Placeholder
-            args_placeholders = tuple(
-                PlaceHolder.insert(
-                    arg, AutogradTensor, owner=sy.local_worker, role=self.role, tracing=True
-                )
-                for arg in args
+            args_placeholders = build_nested_arg(
+                args,
+                lambda x: PlaceHolder.insert(
+                    x, AutogradTensor, owner=sy.local_worker, role=self.role, tracing=True
+                ),
             )
         else:
             # Add Placeholder on top of each arg
-            args = args_placeholders = tuple(
-                PlaceHolder.create_from(arg, owner=sy.local_worker, role=self.role, tracing=True)
-                for arg in args
+            args = args_placeholders = build_nested_arg(
+                args,
+                lambda x: PlaceHolder.create_from(
+                    x, owner=sy.local_worker, role=self.role, tracing=True
+                ),
             )
 
         # Add state to args if needed
@@ -262,6 +281,7 @@ def copy(self):
             id=sy.ID_PROVIDER.pop(),
             owner=self.owner,
             tags=self.tags,
+            input_types=self.input_types,
             description=self.description,
         )
 
@@ -321,6 +341,7 @@ def __call__(self, *args):
                 args = (*args, self.state)
             return self.forward(*args)
         else:
+            self.input_types.input_check(self, args)
             result = self.role.execute(args)
             if len(result) == 1:
                 return result[0]
@@ -539,6 +560,7 @@ def simplify(worker: AbstractWorker, plan: "Plan") -> tuple:
             sy.serde.msgpack.serde._simplify(worker, plan.tags),
             sy.serde.msgpack.serde._simplify(worker, plan.description),
             sy.serde.msgpack.serde._simplify(worker, plan.torchscript),
+            sy.serde.msgpack.serde._simplify(worker, plan.input_types),
         )
 
     @staticmethod
@@ -550,7 +572,7 @@ def detail(worker: AbstractWorker, plan_tuple: tuple) -> "Plan":
         Returns:
             plan: a Plan object
         """
-        (id_, role, include_state, name, tags, description, torchscript) = plan_tuple
+        (id_, role, include_state, name, tags, description, torchscript, input_types,) = plan_tuple
 
         id_ = sy.serde.msgpack.serde._detail(worker, id_)
         role = sy.serde.msgpack.serde._detail(worker, role)
@@ -558,6 +580,7 @@ def detail(worker: AbstractWorker, plan_tuple: tuple) -> "Plan":
         tags = sy.serde.msgpack.serde._detail(worker, tags)
         description = sy.serde.msgpack.serde._detail(worker, description)
         torchscript = sy.serde.msgpack.serde._detail(worker, torchscript)
+        input_types = sy.serde.msgpack.serde._detail(worker, input_types)
 
         plan = sy.Plan(
             role=role,
@@ -568,6 +591,7 @@ def detail(worker: AbstractWorker, plan_tuple: tuple) -> "Plan":
             name=name,
             tags=tags,
             description=description,
+            input_types=input_types,
         )
 
         plan.torchscript = torchscript
@@ -603,6 +627,10 @@ def bufferize(worker: AbstractWorker, plan: "Plan") -> PlanPB:
         if plan.torchscript:
             protobuf_plan.torchscript = plan.torchscript.save_to_buffer()
 
+        if plan.input_types:
+            input_types = sy.serde.protobuf.serde._bufferize(worker, plan.input_types)
+            protobuf_plan.serialized_input.CopyFrom(input_types)
+
         return protobuf_plan
 
     @staticmethod
@@ -621,6 +649,7 @@ def unbufferize(worker: AbstractWorker, protobuf_plan: PlanPB) -> "Plan":
         name = protobuf_plan.name
         tags = set(protobuf_plan.tags) if protobuf_plan.tags else None
         description = protobuf_plan.description if protobuf_plan.description else None
+        input_types = sy.serde.protobuf.serde._unbufferize(worker, protobuf_plan.serialized_input)
 
         plan = Plan(
             role=role,
@@ -631,6 +660,7 @@ def unbufferize(worker: AbstractWorker, protobuf_plan: PlanPB) -> "Plan":
             name=name,
             tags=tags,
             description=description,
+            input_types=input_types,
         )
 
         if protobuf_plan.torchscript:
diff --git a/syft/execution/role.py b/syft/execution/role.py
index 788214dc0da..7cf8033cb33 100644
--- a/syft/execution/role.py
+++ b/syft/execution/role.py
@@ -2,8 +2,7 @@
 from typing import List
 from typing import Tuple
 from typing import Union
-
-import copy
+from typing import Callable
 
 from syft.generic.frameworks import framework_packages
 
@@ -13,8 +12,6 @@
 from syft.execution.placeholder_id import PlaceholderId
 from syft.execution.state import State
 from syft.generic.frameworks.types import FrameworkTensor
-from syft.generic.object import AbstractObject
-from syft.generic.object_storage import ObjectStorage
 from syft.workers.abstract import AbstractWorker
 
 from syft_proto.execution.v1.role_pb2 import Role as RolePB
@@ -54,6 +51,34 @@ def input_placeholders(self):
     def output_placeholders(self):
         return [self.placeholders[id_] for id_ in self.output_placeholder_ids]
 
+    @staticmethod
+    def nested_object_traversal(obj: any, leaf_function: Callable, leaf_type: type):
+        """
+        Class method to iterate through a tree-like structure, where the branching is determined by the elements of list,
+        tuples and dicts, returning the same tree-like structure with a function applied to its leafs.
+
+        Args:
+            obj: The tree-like structure, can be only the root as well.
+            leaf_function: The function to be applied on the leaf nodes of the tree-like structure.
+            leaf_type: On what type on function to apply the function, if the types won't match, the leaf is returned,
+            to apply on all leafs pass any.
+
+        Returns:
+            Same structure as the obj argument, but with the function applied to the leaf elements.
+        """
+        if isinstance(obj, (list, tuple)):
+            result = [Role.nested_object_traversal(elem, leaf_function, leaf_type) for elem in obj]
+            return type(obj)(result)
+        elif isinstance(obj, dict):
+            return {
+                k: Role.nested_object_traversal(v, leaf_function, leaf_type)
+                for k, v in sorted(obj.items())
+            }
+        elif isinstance(obj, leaf_type):
+            return leaf_function(obj)
+        else:
+            return obj
+
     def register_input(self, arg_):
         """ Takes input argument for this role and generate placeholder.
         """
@@ -63,9 +88,14 @@ def register_inputs(self, args_):
         """ Takes input arguments for this role and generate placeholders.
         """
         # TODO Should we be able to rebuild?
-        self.input_placeholder_ids += tuple(
-            self._store_placeholders(arg).value for arg in args_ if isinstance(arg, PlaceHolder)
-        )
+        def traversal_function(obj):
+            if obj.id.value not in self.placeholders:
+                self.placeholders[obj.id.value] = obj
+            self.input_placeholder_ids.append(obj.id.value)
+
+        self.input_placeholder_ids = []
+        Role.nested_object_traversal(args_, traversal_function, PlaceHolder)
+        self.input_placeholder_ids = tuple(self.input_placeholder_ids)
 
     def register_output(self, result):
         """ Takes output tensor for this role and generate placeholder.
@@ -75,11 +105,16 @@ def register_output(self, result):
     def register_outputs(self, results):
         """ Takes output tensors for this role and generate placeholders.
         """
+
+        def traversal_function(obj):
+            if obj.id.value not in self.placeholders:
+                self.placeholders[obj.id.value] = obj
+            self.output_placeholder_ids.append(obj.id.value)
+
         results = (results,) if not isinstance(results, tuple) else results
-        results += tuple(self._store_placeholders(result) for result in results)
-        self.output_placeholder_ids = tuple(
-            result.value for result in results if isinstance(result, PlaceholderId)
-        )
+        self.output_placeholder_ids = []
+        Role.nested_object_traversal(results, traversal_function, PlaceHolder)
+        self.output_placeholder_ids = tuple(self.output_placeholder_ids)
 
     def register_action(self, traced_action, action_type):
         """ Build placeholders and store action.
@@ -119,10 +154,16 @@ def execute(self, args_):
     def _instantiate_inputs(self, args_):
         """ Takes input arguments for this role and generate placeholders.
         """
-        input_placeholders = tuple(
+
+        def traversal_function(obj):
+            placeholder = input_placeholders.pop(0)
+            placeholder.instantiate(obj)
+
+        input_placeholders = [
             self.placeholders[input_id] for input_id in self.input_placeholder_ids
-        )
-        PlaceHolder.instantiate_placeholders(input_placeholders, args_)
+        ]
+
+        Role.nested_object_traversal(args_, traversal_function, FrameworkTensor)
 
     def _execute_action(self, action):
         """ Build placeholders and store action.
@@ -167,31 +208,21 @@ def _store_placeholders(self, obj):
         """
         Replace in an object all FrameworkTensors with Placeholder ids
         """
-        if isinstance(obj, (tuple, list)):
-            r = [self._store_placeholders(o) for o in obj]
-            return type(obj)(r)
-        elif isinstance(obj, dict):
-            return {key: self._store_placeholders(value) for key, value in obj.items()}
-        elif isinstance(obj, PlaceHolder):
+
+        def traversal_function(obj):
             if obj.id.value not in self.placeholders:
                 self.placeholders[obj.id.value] = obj
             return obj.id
-        else:
-            return obj
+
+        return Role.nested_object_traversal(obj, traversal_function, PlaceHolder)
 
     def _fetch_placeholders_from_ids(self, obj):
         """
         Replace in an object all ids with Placeholders
         """
-        if isinstance(obj, (tuple, list)):
-            r = [self._fetch_placeholders_from_ids(o) for o in obj]
-            return type(obj)(r)
-        elif isinstance(obj, dict):
-            return {key: self._fetch_placeholders_from_ids(value) for key, value in obj.items()}
-        elif isinstance(obj, PlaceholderId):
-            return self.placeholders[obj.value]
-        else:
-            return obj
+        return Role.nested_object_traversal(
+            obj, lambda x: self.placeholders[x.value], PlaceholderId
+        )
 
     def copy(self):
         # TODO not the cleanest method ever
@@ -220,16 +251,9 @@ def copy(self):
 
         state = State(state_placeholders)
 
-        def _replace_placeholder_ids(obj):
-            if isinstance(obj, (tuple, list)):
-                r = [_replace_placeholder_ids(o) for o in obj]
-                return type(obj)(r)
-            elif isinstance(obj, dict):
-                return {key: _replace_placeholder_ids(value) for key, value in obj.items()}
-            elif isinstance(obj, PlaceholderId):
-                return PlaceholderId(old_ids_2_new_ids[obj.value])
-            else:
-                return obj
+        _replace_placeholder_ids = lambda obj: Role.nested_object_traversal(
+            obj, lambda x: PlaceholderId(old_ids_2_new_ids[x.value]), PlaceholderId
+        )
 
         new_actions = []
         for action in self.actions:
diff --git a/syft/execution/type_wrapper.py b/syft/execution/type_wrapper.py
new file mode 100644
index 00000000000..2135560f29e
--- /dev/null
+++ b/syft/execution/type_wrapper.py
@@ -0,0 +1,337 @@
+import syft as sy
+from syft.workers.abstract import AbstractWorker
+from warnings import warn
+from typing import Union
+from syft_proto.execution.v1.type_wrapper_pb2 import NestedTypeWrapper as NestedTypeWrapperPB
+from syft_proto.execution.v1.type_wrapper_pb2 import InputTypeDescriptor as InputTypeDescriptorPB
+
+
+class NestedTypeWrapper:
+    """
+        Class for input type serialization and type checking for nested structures.
+    """
+
+    def __init__(self, nested_type=None):
+        self.nested_input_types = NestedTypeWrapper.enumerate_nested_types(nested_type)
+
+    @staticmethod
+    def get_object_identifiers(obj):
+        type_name = type(obj).__name__
+
+        if hasattr(obj, "name"):
+            object_name = obj.name
+        elif hasattr(obj, "id"):
+            object_name = obj.id
+        else:
+            object_name = repr(obj)
+
+        return (type_name, object_name)
+
+    @staticmethod
+    def enumerate_nested_types(input_arg: any) -> Union[list, tuple, dict, type]:
+        """
+            Method to enumerate the input of a function/Plan, including nested types.
+
+            Note: supported nested structures: list, tuple, dict with string keys.
+
+            Params:
+                input_arg: *args of a function or Plan.
+
+            Returns:
+                Union[list, tuple, dict, type]: Nested structured with types instead of objects.
+        """
+        if input_arg is None:
+            return None
+
+        if isinstance(input_arg, (list, tuple)):
+            result = [NestedTypeWrapper.enumerate_nested_types(elem) for elem in input_arg]
+            return tuple(result) if isinstance(input_arg, tuple) else result
+
+        if isinstance(input_arg, dict):
+            serialized_dict = {
+                k: NestedTypeWrapper.enumerate_nested_types(v) for k, v in input_arg.items()
+            }
+            return serialized_dict
+
+        return type(input_arg)
+
+    @staticmethod
+    def raise_typecheck_warn(typechecked_object: any, build: str, call: str, path: str) -> None:
+        """
+            Function to raise a type error if two types differ.
+
+            Params:
+                obj_type: the type of the object returned by calling .__name__ on it.
+                obj_name: the name/id of the object.
+                build: the build/reference argument type.
+                call: the called argument type.
+                path: the nested path to reach that obj.
+
+            Returns:
+                 None
+        """
+        type_name, obj_name = NestedTypeWrapper.get_object_identifiers(typechecked_object)
+        raise TypeError(
+            f"{type_name} {obj_name} {path} has type {build}, while being built with type {call}.",
+        )
+
+    @staticmethod
+    def raise_missmatch_err(typechecked_object: any, build: int, call: int, path: str) -> None:
+        """
+            Function to raise an error if two nested structures differ in length.
+
+            Params:
+                obj_type: the type of the object returned by calling .__name__ on it.
+                obj_name: the name/id of the object.
+                build: the build/reference argument length.
+                call: the called argument length.
+                path: the nested path to reach that obj.
+
+            Returns:
+                 None
+        """
+        type_name, obj_name = NestedTypeWrapper.get_object_identifiers(typechecked_object)
+        raise TypeError(
+            f"{type_name} {obj_name} {path} has length {call}, while being build with length {build}."
+        )
+
+    @staticmethod
+    def raise_wrong_number_arguments_err(typechecked_object: any, build: int, call: int) -> None:
+        """
+            Function to raise an error if the build/reference function has a different number of arguments.
+
+            Params:
+                obj_type: the type of the object returned by calling .__name__ on it.
+                obj_name: the name/id of the object.
+                build: the build/reference input length.
+                call: the called input length.
+
+            Returns:
+                 None
+        """
+        type_name, obj_name = NestedTypeWrapper.get_object_identifiers(typechecked_object)
+        raise TypeError(f"{type_name} {obj_name} requires {build} arguments, received {call}.")
+
+    @staticmethod
+    def raise_key_missing_err(typechecked_object: any, key: any, path: str) -> None:
+        """
+            Function to raise an error if the build/reference function has a different number of arguments.
+
+            Params:
+                obj_type: the type of the object returned by calling .__name__ on type(obj).
+                obj_name: the name/id of the object.
+                key: the key that is missing from the called dict.
+                path: the nested path to reach that obj.
+
+            Returns:
+                 None
+        """
+        type_name, obj_name = NestedTypeWrapper.get_object_identifiers(typechecked_object)
+        raise KeyError(
+            f"{type_name} {obj_name} {path} does not provide the key {key}, while being build with that key."
+        )
+
+    def input_check(self, typechecked_object: any, args: list) -> None:
+        """
+            Method for input validation by comparing the serialized build input with the
+            current call input, following the following steps:
+                1. Input length validation - checking that build and call inputs match on length.
+                2. Verify the following nested structures: list, tuple, dict recursively. Lengths must match when
+                comparing two nested lists, tuples or dicts. If they differ, an error will be raised.
+                3. If we hit an object for which we don't support nesting, we compare types between call input and
+                build input. If they differ, a warning will be raised.
+                4. Dicts on the same nesting level on build and call input must have the same keys. If they differ, an
+                error will be raised.
+
+            Params:
+                obj_type: the type of the object returned by calling .__name__ on type(obj).
+                obj_name: the name/id of the object
+                args: the arguments to be compared with the reference/build one.
+
+            Returns:
+                None
+        """
+
+        def check_type_nested_structure(
+            typechecked_object,
+            build_arg_nested_type: Union[list, tuple, dict, type],
+            call_arg_nested_obj: any,
+            path: str,
+        ) -> None:
+            """
+                Recursive method to compare the nested input argument and the nested build argument.
+
+                Params:
+                    build_arg_nested_type: Can be either a nested element (list, tuple, dict) or a type.
+                    call_arg_nested_obj: Can be either a nested element (list, tuple, dict) or an object.
+
+                Returns:
+                    None
+            """
+
+            iterable_supported_list = (list, tuple, dict)
+
+            if type(call_arg_nested_obj) not in iterable_supported_list:
+                if not isinstance(call_arg_nested_obj, build_arg_nested_type):
+                    NestedTypeWrapper.raise_typecheck_warn(
+                        typechecked_object,
+                        build_arg_nested_type.__name__,
+                        type(call_arg_nested_obj).__name__,
+                        path,
+                    )
+                return
+
+            if type(build_arg_nested_type) != type(call_arg_nested_obj):
+                NestedTypeWrapper.raise_typecheck_warn(
+                    typechecked_object,
+                    type(build_arg_nested_type).__name__,
+                    type(call_arg_nested_obj).__name__,
+                    path,
+                )
+                return
+
+            if isinstance(build_arg_nested_type, (list, tuple)):
+                if len(build_arg_nested_type) != len(call_arg_nested_obj):
+                    NestedTypeWrapper.raise_missmatch_err(
+                        typechecked_object,
+                        len(build_arg_nested_type),
+                        len(call_arg_nested_obj),
+                        path,
+                    )
+
+                for idx in range(len(build_arg_nested_type)):
+                    check_type_nested_structure(
+                        typechecked_object,
+                        build_arg_nested_type[idx],
+                        call_arg_nested_obj[idx],
+                        f"element {idx} of " + path,
+                    )
+
+            if isinstance(build_arg_nested_type, dict):
+                if len(build_arg_nested_type) != len(call_arg_nested_obj):
+                    NestedTypeWrapper.raise_missmatch_err(
+                        typechecked_object,
+                        len(build_arg_nested_type),
+                        len(call_arg_nested_obj),
+                        path,
+                    )
+
+                for key in build_arg_nested_type.keys():
+                    if key in call_arg_nested_obj:
+                        check_type_nested_structure(
+                            typechecked_object,
+                            build_arg_nested_type[key],
+                            call_arg_nested_obj[key],
+                            f"key {key} of " + path,
+                        )
+                    else:
+                        NestedTypeWrapper.raise_key_missing_err(typechecked_object, key, path)
+
+        if len(args) != len(self.nested_input_types):
+            NestedTypeWrapper.raise_wrong_number_arguments_err(
+                typechecked_object, len(self.nested_input_types), len(args)
+            )
+
+        for idx in range(len(args)):
+            check_type_nested_structure(
+                typechecked_object,
+                self.nested_input_types[idx],
+                args[idx],
+                f"element {idx} of input",
+            )
+
+    @staticmethod
+    def simplify(worker: AbstractWorker, nested_type_wrapper: "NestedTypeWrapper") -> list:
+        return sy.serde.msgpack.serde._simplify(worker, nested_type_wrapper.nested_input_types)
+
+    @staticmethod
+    def detail(worker: AbstractWorker, simplified_nested_type: list) -> "NestedTypeWrapper":
+        nested_type_wrapper = sy.serde.msgpack.serde._detail(worker, simplified_nested_type)
+        result = NestedTypeWrapper()
+        result.nested_input_types = nested_type_wrapper
+        return result
+
+    @staticmethod
+    def bufferize(
+        worker: AbstractWorker, nested_type_wrapper: "NestedTypeWrapper"
+    ) -> NestedTypeWrapperPB:
+        def bufferize_nested_structure(worker: AbstractWorker, obj: any) -> NestedTypeWrapperPB:
+            nested_type_pb = NestedTypeWrapperPB()
+
+            if isinstance(obj, list):
+                container = NestedTypeWrapperPB.TypeContainer()
+                proto_list = NestedTypeWrapperPB.TypeList()
+
+                for elem in obj:
+                    proto_list.nested_types.append(bufferize_nested_structure(worker, elem))
+
+                container.nested_type_list.CopyFrom(proto_list)
+                nested_type_pb.nested_types.CopyFrom(container)
+
+            if isinstance(obj, tuple):
+                container = NestedTypeWrapperPB.TypeContainer()
+                proto_tuple = NestedTypeWrapperPB.TypeTuple()
+
+                for elem in obj:
+                    proto_tuple.nested_types.append(bufferize_nested_structure(worker, elem))
+
+                container.nested_type_tuple.CopyFrom(proto_tuple)
+                nested_type_pb.nested_types.CopyFrom(container)
+
+            if isinstance(obj, dict):
+                container = NestedTypeWrapperPB.TypeContainer()
+                proto_map = NestedTypeWrapperPB.TypeMap()
+
+                for k, v in obj.items():
+                    proto_map.nested_types[k].CopyFrom(bufferize_nested_structure(worker, v))
+
+                container.nested_type_dict.CopyFrom(proto_map)
+                nested_type_pb.nested_types.CopyFrom(container)
+
+            if isinstance(obj, type):
+                container = NestedTypeWrapperPB.TypeContainer()
+                typePB = InputTypeDescriptorPB()
+                module_path = obj.__module__
+                full_path_type = module_path + "." + obj.__name__
+                typePB.type_name = full_path_type
+                container.nested_type.CopyFrom(typePB)
+                nested_type_pb.nested_types.CopyFrom(container)
+            return nested_type_pb
+
+        result = bufferize_nested_structure(worker, nested_type_wrapper.nested_input_types)
+        return result
+
+    @staticmethod
+    def unbufferize(worker: AbstractWorker, message):
+        def unbufferize_nested_structure(worker, message):
+            container = None
+            if message.nested_types.HasField("nested_type"):
+                return sy.serde.protobuf.serde._unbufferize(
+                    worker, message.nested_types.nested_type
+                )
+
+            if message.nested_types.HasField("nested_type_list"):
+                container = []
+                for obj in message.nested_types.nested_type_list.nested_types:
+                    container.append(unbufferize_nested_structure(worker, obj))
+
+            if message.nested_types.HasField("nested_type_tuple"):
+                container = []
+                for obj in message.nested_types.nested_type_tuple.nested_types:
+                    container.append(unbufferize_nested_structure(worker, obj))
+                container = tuple(container)
+
+            if message.nested_types.HasField("nested_type_dict"):
+                container = {}
+                for k, v in message.nested_types.nested_type_dict.nested_types.items():
+                    container[k] = unbufferize_nested_structure(worker, v)
+
+            return container
+
+        result = unbufferize_nested_structure(worker, message)
+        wrapper = NestedTypeWrapper()
+        wrapper.nested_input_types = result
+        return wrapper
+
+    def __iter__(self):
+        return iter(self.nested_input_types)
diff --git a/syft/serde/msgpack/native_serde.py b/syft/serde/msgpack/native_serde.py
index cc5249c05f7..f865a5ad46e 100644
--- a/syft/serde/msgpack/native_serde.py
+++ b/syft/serde/msgpack/native_serde.py
@@ -7,7 +7,8 @@
 from typing import Dict
 from typing import Union
 from typing import Tuple
-
+import warnings
+import pydoc
 
 import numpy
 
@@ -364,6 +365,54 @@ def _detail_ndarray(
     return res
 
 
+def _simplify_type(worker: AbstractWorker, obj_type: type) -> bytes:
+    """
+    This function gets an type object and returns its representation as bytes.
+
+    Args:
+        obj_type (s.g builtins.str, builtins.int, torch.tensor): a type
+
+    Returns:
+        str: a string in utf-8 encoding that encodes the path in the module + the actual class that generates the type.
+
+    Examples:
+          str_type_representation = _simplify_type(worker, type("i'm a string"))
+    """
+    module_path = obj_type.__module__
+    full_path_type = module_path + "." + obj_type.__name__
+    return full_path_type.encode("utf-8")
+
+
+def _detail_type(worker: AbstractWorker, type_repr: bytes) -> type:
+    """
+    This function gets the byte representation of a type and its path in the module as a string, decodes the string
+    and locates the type in the module, returning the type object.
+
+    Args:
+        type_repr: bytes that encode the path of a type/class in a module
+
+    Returns:
+        type: the type of an object (e.g: builtins.str, builtins.int).
+
+    Warning: if pydoc can't locate the type in the current process, might mean that the file layout is different between
+    sender and receiver.
+
+    TODO:
+        As syft-protobuf grows in type support, we should change the type serialization by using those types, enabling cross
+        language typechecking/type validation.
+    """
+
+    string_representation = type_repr.decode("utf-8")
+    result = pydoc.locate(string_representation)
+    if result is None:
+        warnings.warn(
+            f"{string_representation} can't be located in the current process, the layout of the modules has been changed.",
+            Warning,
+        )
+        return object
+    return result
+
+
 def _simplify_numpy_number(
     worker: AbstractWorker, numpy_nb: Union[numpy.int32, numpy.int64, numpy.float32, numpy.float64]
 ) -> Tuple[bin, Tuple]:
@@ -425,6 +474,7 @@ def _detail_numpy_number(
         range: (_simplify_range, _detail_range),
         set: (_simplify_collection, _detail_collection_set),
         slice: (_simplify_slice, _detail_slice),
+        type: (_simplify_type, _detail_type),
         str: (_simplify_str, _detail_str),
         tuple: (_simplify_collection, _detail_collection_tuple),
         type(Ellipsis): (_simplify_ellipsis, _detail_ellipsis),
diff --git a/syft/serde/msgpack/serde.py b/syft/serde/msgpack/serde.py
index f81b90ce349..a0909f365e4 100644
--- a/syft/serde/msgpack/serde.py
+++ b/syft/serde/msgpack/serde.py
@@ -59,6 +59,8 @@
 from syft.execution.state import State
 from syft.execution.computation import ComputationAction
 from syft.execution.communication import CommunicationAction
+from syft.execution.protocol import Protocol
+from syft.execution.type_wrapper import NestedTypeWrapper
 from syft.messaging.message import TensorCommandMessage
 from syft.messaging.message import ObjectMessage
 from syft.messaging.message import ObjectRequestMessage
@@ -116,6 +118,7 @@
     Role,
     ObjectPointer,
     Plan,
+    NestedTypeWrapper,
     Protocol,
     State,
     ComputationAction,
diff --git a/syft/serde/protobuf/native_serde.py b/syft/serde/protobuf/native_serde.py
index e930ade9fb5..351abfedc29 100644
--- a/syft/serde/protobuf/native_serde.py
+++ b/syft/serde/protobuf/native_serde.py
@@ -4,11 +4,15 @@
 something here that isn't for `None`, think twice and either
 use an existing sub-class of Message or add a new one.
 """
-
+import pydoc
+import warnings
 from collections import OrderedDict
+
 from google.protobuf.empty_pb2 import Empty
 from syft.workers.abstract import AbstractWorker
 
+from syft_proto.execution.v1.type_wrapper_pb2 import InputTypeDescriptor as InputTypeDescriptorPB
+
 
 def _bufferize_none(worker: AbstractWorker, obj: "type(None)") -> "Empty":
     """
@@ -36,5 +40,60 @@ def _unbufferize_none(worker: AbstractWorker, obj: "Empty") -> "type(None)":
     return None
 
 
+def _bufferize_type(worker: AbstractWorker, obj) -> InputTypeDescriptorPB:
+    """
+    This function gets the type object and returns the ClassType Protobuf message containing the string with the path
+    of that that and the actual type..
+
+    Args:
+        obj_type (s.g builtins.str, builtins.int, torch.tensor): a type
+
+    Returns:
+        ClassTypePB: the Protobuf message type containg the path where to find the type + type.
+
+    Examples:
+          str_type_representation = _bufferize_type(worker, type("i'm a string"))
+    """
+
+    proto_type = InputTypeDescriptorPB()
+
+    if isinstance(obj, type):
+        module_path = obj.__module__
+        full_path_type = module_path + "." + obj.__name__
+        proto_type.type_name = full_path_type
+
+    return proto_type
+
+
+def _unbufferize_type(worker: AbstractWorker, class_type_msg: InputTypeDescriptorPB):
+    """
+    This function receives the ClassType Protobuf message containing the string with the path + type, decodes the string
+    and locates the type in a module, returning the type object.
+
+    Args:
+        class_type_msg: message encoding the type.
+
+    Returns:
+        type: the type of an object (e.g: builtins.str, builtins.int).
+
+    Warning: if pydoc can't locate the type in the current process, might mean that the file layout is different between
+    sender and receiver.
+
+    TODO:
+        As syft-protobuf grows in type support, we should change the type serialization by using those types, enabling cross
+        language typechecking/type validation.
+    """
+    result = pydoc.locate(class_type_msg.type_name)
+    if result is None:
+        warnings.warn(
+            f"{class_type_msg.type_name} can't be located in the current process, the layout of the modules has been changed.",
+            Warning,
+        )
+        return object
+    return result
+
+
 # Maps a type to its bufferizer and unbufferizer functions
-MAP_NATIVE_PROTOBUF_TRANSLATORS = OrderedDict({type(None): (_bufferize_none, _unbufferize_none)})
+MAP_NATIVE_PROTOBUF_TRANSLATORS = OrderedDict(
+    {type(None): (_bufferize_none, _unbufferize_none), type: (_bufferize_type, _unbufferize_type)}
+)
diff --git a/syft/serde/protobuf/proto.py b/syft/serde/protobuf/proto.py
index ce98af249c2..35afe6ffab0 100644
--- a/syft/serde/protobuf/proto.py
+++ b/syft/serde/protobuf/proto.py
@@ -13,6 +13,7 @@
 from syft.execution.placeholder import PlaceHolder
 from syft.execution.placeholder_id import PlaceholderId
 from syft.execution.plan import Plan
+from syft.execution.plan import NestedTypeWrapper
 from syft.execution.protocol import Protocol
 from syft.execution.role import Role
 from syft.execution.state import State
@@ -27,6 +28,8 @@
 )
 from syft_proto.execution.v1.computation_action_pb2 import ComputationAction as ComputationActionPB
 from syft_proto.execution.v1.plan_pb2 import Plan as PlanPB
+from syft_proto.execution.v1.type_wrapper_pb2 import InputTypeDescriptor as InputTypeDescriptorPB
+from syft_proto.execution.v1.type_wrapper_pb2 import NestedTypeWrapper as NestedTypeWrapperPB
 from syft_proto.execution.v1.protocol_pb2 import Protocol as ProtocolPB
 from syft_proto.execution.v1.role_pb2 import Role as RolePB
 from syft_proto.execution.v1.state_pb2 import State as StatePB
@@ -49,6 +52,7 @@
 
 MAP_PYTHON_TO_PROTOBUF_CLASSES = {
     type(None): Empty,
+    type: InputTypeDescriptorPB,
     # Torch types
     torch.Tensor: TorchTensorPB,
     torch.device: DevicePB,
@@ -66,6 +70,7 @@
     ComputationAction: ComputationActionPB,
     PlaceHolder: PlaceholderPB,
     Plan: PlanPB,
+    NestedTypeWrapper: NestedTypeWrapperPB,
     PointerTensor: PointerTensorPB,
     Protocol: ProtocolPB,
     Role: RolePB,
diff --git a/syft/serde/protobuf/serde.py b/syft/serde/protobuf/serde.py
index 8de81bfaff3..c26f9287522 100644
--- a/syft/serde/protobuf/serde.py
+++ b/syft/serde/protobuf/serde.py
@@ -13,6 +13,7 @@
 from syft.messaging.message import ObjectMessage
 from syft.messaging.message import TensorCommandMessage
 from syft.execution.plan import Plan
+from syft.execution.plan import NestedTypeWrapper
 from syft.execution.protocol import Protocol
 from syft.execution.role import Role
 from syft.execution.state import State
@@ -51,6 +52,7 @@
     PlaceholderId,
     PlaceHolder,
     Plan,
+    NestedTypeWrapper,
     PointerTensor,
     Protocol,
     Role,
diff --git a/test/execution/test_plan.py b/test/execution/test_plan.py
index 1a8816be6e4..9fa2dd59ef3 100644
--- a/test/execution/test_plan.py
+++ b/test/execution/test_plan.py
@@ -618,9 +618,47 @@ def plan_abs(data):
     assert len(pointers) == 2
 
 
+def test_plan_input_serialization(hook):
+    @sy.func2plan()
+    def plan_serialized_input_1(x, y, z, t):
+        return x
+
+    plan_serialized_input_1.build(
+        (th.tensor([1.0, -2.0]), th.tensor([1, 2])),
+        {
+            "k1": [5, (True, False)],
+            "k2": {
+                "kk1": [th.tensor([5, 7]), th.tensor([5, 7])],
+                "kk2": [True, (th.tensor([9, 10]),)],
+            },
+            "k3": th.tensor([8]),
+        },
+        th.tensor([11, 12]),
+        (1, (2, (3, (4, [5, 6])))),
+    )
+
+    reference_serialized_input_1 = (
+        (type(th.tensor([1.0, -2.0])), type(th.tensor([1, 2]))),
+        {
+            "k1": [type(5), (type(True), type(False))],
+            "k2": {
+                "kk1": [type(th.tensor([5, 7])), type(th.tensor([5, 7]))],
+                "kk2": [type(True), (type(th.tensor([9, 10])),)],
+            },
+            "k3": type(th.tensor([8])),
+        },
+        type(th.tensor([11, 12])),
+        (type(1), (type(2), (type(3), (type(4), [type(5), type(6)])))),
+    )
+
+    serialized_input_1 = plan_serialized_input_1.input_types.nested_input_types
+
+    assert reference_serialized_input_1 == serialized_input_1
+
+
 def test_plan_input_usage(hook):
-    x11 = th.tensor([-1, 2.0]).tag("input_data")
-    x12 = th.tensor([1, -2.0]).tag("input_data2")
+    x11 = th.tensor([-1, 2.0])
+    x12 = th.tensor([1, -2.0])
 
     device_1 = sy.VirtualWorker(hook, id="test_dev_1", data=(x11, x12))
 
@@ -632,8 +670,8 @@ def plan_test_1(x, y):
     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]
+    pointer_to_data_1 = x11.send(device_1)
+    pointer_to_data_2 = x12.send(device_1)
 
     plan_test_1.build(th.tensor([1.0, -2.0]), th.tensor([1, 2]))
     pointer_plan = plan_test_1.send(device_1)
@@ -641,11 +679,231 @@ def plan_test_2(x, y):
     result = pointer_to_result.get()
     assert (result == x11).all()
 
+    pointer_to_data_1 = x11.send(device_1)
+    pointer_to_data_2 = x12.send(device_1)
+
     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
+    assert (result == x12).all()
+
+
+def test_plan_wrong_number_of_parameters(hook):
+    x11 = th.tensor([-1, 2.0])
+    x12 = th.tensor([1, -2.0])
+
+    device_1 = sy.VirtualWorker(hook, id="test_dev_1", data=(x11, x12))
+
+    @sy.func2plan()
+    def plan_test(x, y, z, t):
+        return x
+
+    pointer_to_data_1 = x11.send(device_1)
+    pointer_to_data_2 = x12.send(device_1)
+
+    dummy_input_list = [th.tensor([1, -2]), th.tensor([1, 2]), 2, False]
+    input_list = [pointer_to_data_1, pointer_to_data_2, 5, True]
+    corect_number_of_params = 4
+    plan_test.build(*dummy_input_list)
+    pointer_plan = plan_test.send(device_1)
+
+    for i in range(len(input_list) + 1):
+        if i == corect_number_of_params:
+            pointer_plan(*input_list[:i])
+        else:
+            with pytest.raises(TypeError) as e:
+                pointer_plan(*input_list[:i])
+
+            assert f"Plan plan_test requires {len(input_list)} arguments, received {i}." == str(
+                e.value
+            )
+
+
+def test_plan_list(hook):
+    x11 = th.tensor([-1, 2.0])
+    x12 = th.tensor([1, -2.0])
+
+    @sy.func2plan()
+    def plan_list(data, x):
+        y = data[0] + data[1]
+        z = data[0] + x
+        return y + z
+
+    device_1 = sy.VirtualWorker(hook, id="test_plan_list", data=(x11, x12))
+
+    plan_list.build([th.tensor([1, 2]), th.tensor([2, 3])], th.tensor([0, 0]))
+    pointer_to_plan = plan_list.send(device_1)
+    pointer_to_data_1 = x11.send(device_1)
+    pointer_to_data_2 = x12.send(device_1)
+    result = pointer_to_plan([pointer_to_data_1, pointer_to_data_2], th.tensor([1, 1]))
+    assert (result.get() == th.tensor([0, 3])).all()
+
+
+def test_plan_tuple(hook):
+    x11 = th.tensor([-1, 2.0])
+    x12 = th.tensor([1, -2.0])
+
+    @sy.func2plan()
+    def plan_tuple(data, x):
+        y = data[0] + data[1]
+        z = data[0] + x
+        return y + z
+
+    device_1 = sy.VirtualWorker(hook, id="test_plan_tuple", data=(x11, x12))
+
+    plan_tuple.build((th.tensor([1, 2]), th.tensor([2, 3])), th.tensor([0, 0]))
+    pointer_to_plan = plan_tuple.send(device_1)
+
+    pointer_to_data_1 = x11.send(device_1)
+    pointer_to_data_2 = x12.send(device_1)
+
+    result = pointer_to_plan((pointer_to_data_1, pointer_to_data_2), th.tensor([1, 1]))
+    assert (result.get() == th.tensor([0, 3])).all()
+
+
+def test_plan_dict(hook):
+    x11 = th.tensor([-1, 2.0])
+    x12 = th.tensor([1, -2.0])
+
+    @sy.func2plan()
+    def plan_dict(data, x):
+        y = data["input1"] + data["input2"]
+        z = data["input1"] + x
+        return y + z
+
+    device_1 = sy.VirtualWorker(hook, id="test_plan_dict", data=(x11, x12))
+    plan_dict.build({"input1": th.tensor([1, 2]), "input2": th.tensor([2, 3])}, th.tensor([0, 0]))
+    pointer_to_plan = plan_dict.send(device_1)
+    pointer_to_data_1 = x11.send(device_1)
+    pointer_to_data_2 = x12.send(device_1)
+    result = pointer_to_plan(
+        {"input1": pointer_to_data_1, "input2": pointer_to_data_2}, th.tensor([1, 1])
+    )
+    assert (result.get() == th.tensor([0, 3])).all()
+
+
+def test_plan_nested_structures(hook):
+    x1 = th.tensor([-1, 2.0])
+    x2 = th.tensor([1, -2.0])
+    x3 = th.tensor([2, -1])
+    x4 = th.tensor([2, 1])
+
+    @sy.func2plan()
+    def plan_nested(dict):
+        """
+        dict:
+            "tensors":
+                "test1": [tensor(-1, 2), tensor(1, -2), [tensor(-1, 2), tensor(2, 1)]]
+                "test2": tensor(2, 1)
+            "tensor_list": [tensor(-1, 2), tensor(2, -1)]
+        """
+        x = dict["tensors"]["test1"][0]
+        y = dict["tensors"]["test1"][1]
+        z = dict["tensors"]["test2"]
+        t = dict["tensor_list"][0]
+        return x + y + z + t
+
+    dummy_build = {
+        "tensors": {
+            "test1": [th.tensor([0, 0]), th.tensor([0, 0]), [th.tensor([0, 0]), th.tensor([0, 0])]],
+            "test2": th.tensor([0, 0]),
+        },
+        "tensor_list": [th.tensor([-1, 2]), th.tensor([1, -2])],
+    }
+
+    device_1 = sy.VirtualWorker(hook, id="test_nested_structure", data=(x1, x2, x3, x4))
+    plan_nested.build(dummy_build)
+
+    pointer_to_data_1 = x1.send(device_1)
+    pointer_to_data_2 = x2.send(device_1)
+    pointer_to_data_3 = x3.send(device_1)
+    pointer_to_data_4 = x4.send(device_1)
+
+    call_build = {
+        "tensors": {
+            "test1": [pointer_to_data_1, pointer_to_data_2, [pointer_to_data_1, pointer_to_data_4]],
+            "test2": pointer_to_data_4,
+        },
+        "tensor_list": [pointer_to_data_1, pointer_to_data_3],
+    }
+    pointer_to_plan = plan_nested.send(device_1)
+    result = pointer_to_plan(call_build)
+    assert (result.get() == th.tensor([1, 3])).all()
+
+
+def test_plan_type_error(hook):
+    x1 = th.tensor([-1, 2.0])
+    x2 = th.tensor([1, -2.0])
+
+    @sy.func2plan()
+    def plan_type_err(dic):
+        return dic["k1"]["kk2"]
+
+    dummy_build = {
+        "k1": {
+            "kk1": [(th.tensor([0, 0]), 1), [th.tensor([0, 0]), 2.5]],
+            "kk2": th.tensor([0, 0]),
+        },
+        "k2": [th.tensor([-1, 2]), "dummy_str"],
+    }
+
+    device_1 = sy.VirtualWorker(hook, id="test_nested_structure", data=(x1, x2))
+    plan_type_err.build(dummy_build)
+
+    pointer_to_data_1 = x1.send(device_1)
+    pointer_to_data_2 = x2.send(device_1)
+
+    call_build = {
+        "k1": {"kk1": [(pointer_to_data_1, 1.5), [pointer_to_data_1, True]], "kk2": "warn",},
+        "k2": [pointer_to_data_1, pointer_to_data_2],
+    }
+
+    pointer_to_plan = plan_type_err.send(device_1)
+    with pytest.raises(TypeError) as e:
+        _ = pointer_to_plan(call_build)
+
+    assert (
+        str(e.value)
+        == "Plan plan_type_err element 1 of element 0 of key kk1 of key k1 of element 0 of input has type int, while being built with type float."
+    )
+
+
+def test_plan_key_error(hook):
+    x1 = th.tensor([-1, 2.0])
+    x2 = th.tensor([1, -2.0])
+
+    @sy.func2plan()
+    def plan_type_warn(dic):
+        return dic["k1"]["kk2"]
+
+    dummy_build = {
+        "k1": {
+            "kk1": [(th.tensor([0, 0]), 1), [th.tensor([0, 0]), 2.5]],
+            "kk2": th.tensor([0, 0]),
+        },
+        "k2": [th.tensor([-1, 2]), "dummy_str"],
+    }
+
+    device_1 = sy.VirtualWorker(hook, id="test_nested_structure", data=(x1, x2))
+    plan_type_warn.build(dummy_build)
+
+    pointer_to_data_1 = x1.send(device_1)
+    pointer_to_data_2 = x2.send(device_1)
+
+    call_build = {
+        "k1": {"kk1_wrong": [(pointer_to_data_1, 1.5), [pointer_to_data_1, True]], "kk2": "warn",},
+        "k2": [pointer_to_data_1, pointer_to_data_2],
+    }
+
+    pointer_to_plan = plan_type_warn.send(device_1)
+    with pytest.raises(KeyError) as e:
+        _ = pointer_to_plan(call_build)
+
+    assert (
+        str(e.value)
+        == "'Plan plan_type_warn key k1 of element 0 of input does not provide the key kk1, while being build with that key.'"
+    )
 
 
 def test_backward_autograd_can_be_traced(hook, workers):
diff --git a/test/serde/msgpack/test_msgpack_serde_full.py b/test/serde/msgpack/test_msgpack_serde_full.py
index dfcba4b4e41..ef6c468c0aa 100644
--- a/test/serde/msgpack/test_msgpack_serde_full.py
+++ b/test/serde/msgpack/test_msgpack_serde_full.py
@@ -24,6 +24,7 @@
 samples[str] = make_str
 samples[range] = make_range
 samples[type(Ellipsis)] = make_ellipsis
+samples[type] = make_type
 
 # Numpy
 samples[numpy.float32] = partial(make_numpy_number, numpy.float32)
@@ -88,6 +89,7 @@
 samples[syft.messaging.message.TensorCommandMessage] = make_command_message
 samples[syft.messaging.message.WorkerCommandMessage] = make_workercommandmessage
 
+samples[syft.execution.plan.NestedTypeWrapper] = make_nested_type_wrapper
 samples[syft.workers.base.BaseWorker] = make_baseworker
 
 
diff --git a/test/serde/protobuf/test_protobuf_serde_full.py b/test/serde/protobuf/test_protobuf_serde_full.py
index ee8f9d4f2f9..61aa35c0713 100644
--- a/test/serde/protobuf/test_protobuf_serde_full.py
+++ b/test/serde/protobuf/test_protobuf_serde_full.py
@@ -15,6 +15,7 @@
 
 # Native
 samples[type(None)] = make_none
+samples[type] = make_type
 
 # PyTorch
 samples[torch.device] = make_torch_device
@@ -37,6 +38,7 @@
 samples[syft.execution.role.Role] = make_role
 samples[syft.execution.state.State] = make_state
 samples[syft.execution.placeholder_id.PlaceholderId] = make_placeholder_id
+samples[syft.execution.plan.NestedTypeWrapper] = make_nested_type_wrapper
 samples[syft.generic.pointers.pointer_tensor.PointerTensor] = make_pointertensor
 
 # Syft Messages
diff --git a/test/serde/serde_helpers.py b/test/serde/serde_helpers.py
index a50f989dc77..7dcbeb11725 100644
--- a/test/serde/serde_helpers.py
+++ b/test/serde/serde_helpers.py
@@ -808,6 +808,7 @@ def compare(detailed, original):
                     msgpack.serde._simplify(
                         kwargs["workers"]["serde_worker"], plan.torchscript
                     ),  # Torchscript
+                    msgpack.serde._simplify(syft.hook.local_worker, plan.input_types),
                 ),
             ),
             "cmp_detailed": compare,
@@ -828,6 +829,7 @@ def compare(detailed, original):
                     msgpack.serde._simplify(
                         kwargs["workers"]["serde_worker"], model_plan.torchscript
                     ),  # Torchscript
+                    msgpack.serde._simplify(syft.hook.local_worker, plan.input_types),
                 ),
             ),
             "cmp_detailed": compare,
@@ -870,6 +872,54 @@ def compare(detailed, original):
     ]
 
 
+def make_type(**kwargs):
+    serialized_type = type("test")
+
+    def compare(detailed, original):
+        assert type(detailed) == type(original)
+        assert detailed == original
+        return True
+
+    return [
+        {
+            "value": serialized_type,
+            "simplified": (msgpack.serde._simplify(syft.hook.local_worker, serialized_type)),
+            "cmp_detailed": compare,
+        }
+    ]
+
+
+def make_nested_type_wrapper(**kwargs):
+    reference_serialized_input = (
+        (type(torch.tensor([1.0, -2.0])), type(torch.tensor([1, 2]))),
+        {
+            "k1": [type(5), (type(True), type(False))],
+            "k2": {
+                "kk1": [type(torch.tensor([5, 7])), type(torch.tensor([5, 7]))],
+                "kk2": [type(True), (type(torch.tensor([9, 10])),)],
+            },
+            "k3": type(torch.tensor([8])),
+        },
+        type(torch.tensor([11, 12])),
+        (type(1), (type(2), (type(3), (type(4), [type(5), type(6)])))),
+    )
+
+    wrapper = syft.execution.plan.NestedTypeWrapper()
+    wrapper.serialized_nested_type = reference_serialized_input
+
+    def compare(detailed, original):
+        assert detailed.nested_input_types == original.nested_input_types
+        return True
+
+    return [
+        {
+            "value": wrapper,
+            "simplified": syft.serde.msgpack.serde._simplify(syft.hook.local_worker, wrapper),
+            "cmp_detailed": compare,
+        }
+    ]
+
+
 # State
 def make_state(**kwargs):
     me = kwargs["workers"]["me"]


Code review comments:
diff_hunk: @@ -904,3 +904,28 @@ def plan_type_warn(dic):
         str(e.value)
         == "'Plan plan_type_warn key k1 of element 0 of input does not provide the key kk1, while being build with that key.'"
     )
+
+
+def test_backward_autograd_can_be_traced(hook, workers):
+    @sy.func2plan(args_shape=[(5, 5)], trace_autograd=True)
+    def autograd_test(X):
+        y = X * 5
+        y = -y.log() / 2
+        y = y.sum()
+        y.backward()
+        return X.grad
+
+    X = th.ones(5, 5, requires_grad=True)
+
+    # Result of torch autograd
+    torch_grads = autograd_test(X)
+
+    # Result of traced backprop
+    autograd_test.forward = None
+    plan_grads = autograd_test(X)
+
+    # (debug out)
+    print("Traced Plan:\n", autograd_test.code)
path: test/execution/test_plan.py
text: @user1:
Maybe remove that