# Buggy code
def _get_hooked_method(cls, tensor_type, method_name):
    """
    Hook a method in order to replace all args/kwargs syft/torch tensors with
    their child attribute if they exist
    If so, forward this method with the new args and new self, get response
    and "rebuild" the torch tensor wrapper upon all tensors found
    If not, just execute the native torch methodn

    Args:
        attr (str): the method to hook
    Return:
        the hooked method
    """

    @wraps(getattr(tensor_type, method_name))
    def overloaded_native_method(self, *args, **kwargs):
        """
        Operate the hooking
        """

        if not hasattr(self, "child"):  # means that it's not a wrapper
            # if self is a natural tensor but the first argument isn't,
            # wrap self with the appropriate type and re-run
            if len(args) > 0 and hasattr(args[0], "child"):
                # if we allow this for PointerTensors it opens the potential
                # that we could accidentally serialize and send a tensor in the
                # arguments
                if not isinstance(args[0].child, PointerTensor):
                    self = type(args[0].child)().on(self, wrap=True)
                    args = [args[0]]
                    return overloaded_native_method(self, *args, **kwargs)

            method = getattr(self, f"native_{method_name}")
            # Run the native function with the new args

            try:
                response = method(*args, **kwargs)

            except BaseException as e:
                # we can make some errors more descriptive with this method
                raise route_method_exception(e, self, args, kwargs)

        else:  # means that there is a wrapper to remove
            try:
                # Replace all torch tensor with their child attribute
                new_self, new_args, new_kwargs = hook_args.unwrap_args_from_method(
                    method_name, self, args, kwargs
                )

            except BaseException as e:  # if there's a type mismatch, try to fix it!
                try:
                    # if the first argument has no child (meaning it's probably raw data),
                    # try wrapping it with the type of self. We have to except PointerTensor
                    # because otherwise it can lead to inadvertently sending data to another
                    # machine
                    if not hasattr(args[0], "child") and not isinstance(
                        self.child, PointerTensor
                    ):
                        # TODO: add check to make sure this isn't getting around a security class

                        _args = list()
                        _args.append(type(self)().on(args[0], wrap=False))
                        for a in args[1:]:
                            _args.append(a)

                        args = _args

                    # Replace all torch tensor with their child attribute
                    new_self, new_args, new_kwargs = hook_args.unwrap_args_from_method(
                        method_name, self, args, kwargs
                    )
                except BaseException as e:
                    # we can make some errors more descriptive with this method
                    raise route_method_exception(e, self, args, kwargs)

            # Send the new command to the appropriate class and get the response
            method = getattr(new_self, method_name)
            response = method(*new_args, **new_kwargs)

            # For inplace methods, just directly return self
            if syft.framework.is_inplace_method(method_name):
                return self

            # if object is a pointer of pointer, set register to False
            if isinstance(self.child, PointerTensor):
                wrap_args = {"register": False}
            else:
                wrap_args = {}
            # Put back the wrappers where needed
            response = hook_args.hook_response(
                method_name,
                response,
                wrap_type=type(self),
                new_self=self,
                wrap_args=wrap_args,
            )

        return response

    return overloaded_native_method


# Traceback (PureTorchTensorFoundError)
WARNING: Logging before flag parsing goes to stderr.
W0813 08:46:34.190028 140520788572032 secure_random.py:26] Falling back to insecure randomness since the required custom op could not be found for the installed version of TensorFlow. Fix this by compiling custom ops. Missing file was '/usr/local/lib/python3.6/dist-packages/tf_encrypted/operations/secure_random/secure_random_module_tf_1.14.0.so'
W0813 08:46:34.211576 140520788572032 deprecation_wrapper.py:119] From /usr/local/lib/python3.6/dist-packages/tf_encrypted/session.py:26: The name tf.Session is deprecated. Please use tf.compat.v1.Session instead.

0it [00:00, ?it/s]

Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ../data/cifar-10-python.tar.gz

170500096it [00:04, 40609204.86it/s]

---------------------------------------------------------------------------

PureTorchTensorFoundError                 Traceback (most recent call last)

/content/PySyft/syft/frameworks/torch/tensors/interpreters/native.py in handle_func_command(cls, command)
300             new_args, new_kwargs, new_type, args_type = syft.frameworks.torch.hook_args.unwrap_args_from_function(
--> 301                 cmd, args, kwargs, return_args_type=True
302             )

16 frames

PureTorchTensorFoundError:


During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)

/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in batch_norm(input, running_mean, running_var, weight, bias, training, momentum, eps)
1695     return torch.batch_norm(
1696         input, weight, bias, running_mean, running_var,
-> 1697         training, momentum, eps, torch.backends.cudnn.enabled
1698     )
1699

RuntimeError: running_mean should contain 3 elements not 32

# Fixed code
def _get_hooked_method(cls, tensor_type, method_name):
    """
    Hook a method in order to replace all args/kwargs syft/torch tensors with
    their child attribute if they exist
    If so, forward this method with the new args and new self, get response
    and "rebuild" the torch tensor wrapper upon all tensors found
    If not, just execute the native torch methodn

    Args:
        attr (str): the method to hook
    Return:
        the hooked method
    """

    @wraps(getattr(tensor_type, method_name))
    def overloaded_native_method(self, *args, **kwargs):
        """
        Operate the hooking
        """

        if not hasattr(self, "child"):  # means that it's not a wrapper
            # if self is a natural tensor but the first argument isn't,
            # wrap self with the appropriate type and re-run
            if len(args) > 0 and hasattr(args[0], "child"):
                # if we allow this for PointerTensors it opens the potential
                # that we could accidentally serialize and send a tensor in the
                # arguments
                if not isinstance(args[0].child, PointerTensor):
                    self = type(args[0].child)().on(self, wrap=True)
                    args = [args[0]]
                    return overloaded_native_method(self, *args, **kwargs)

            method = getattr(self, f"native_{method_name}")
            # Run the native function with the new args

            try:
                response = method(*args, **kwargs)

            except BaseException as e:
                # we can make some errors more descriptive with this method
                raise route_method_exception(e, self, args, kwargs)

        else:  # means that there is a wrapper to remove
            try:
                # Replace all torch tensor with their child attribute
                new_self, new_args, new_kwargs = hook_args.unwrap_args_from_method(
                    method_name, self, args, kwargs
                )

            except BaseException as e:  # if there's a type mismatch, try to fix it!
                try:
                    # if the first argument has no child (meaning it's probably raw data),
                    # try wrapping it with the type of self. We have to except PointerTensor
                    # because otherwise it can lead to inadvertently sending data to another
                    # machine
                    if not hasattr(args[0], "child") and not isinstance(
                        self.child, PointerTensor
                    ):
                        # TODO: add check to make sure this isn't getting around a security class

                        _args = list()
                        _args.append(type(self)().on(args[0], wrap=False))
                        for a in args[1:]:
                            _args.append(a)

                        args = _args

                    # Replace all torch tensor with their child attribute
                    new_self, new_args, new_kwargs = hook_args.unwrap_args_from_method(
                        method_name, self, args, kwargs
                    )
                except BaseException as e:
                    # we can make some errors more descriptive with this method
                    raise route_method_exception(e, self, args, kwargs)

            # Send the new command to the appropriate class and get the response
            method = getattr(new_self, method_name)
            response = method(*new_args, **new_kwargs)

            # For inplace methods, just directly return self
            if syft.framework.is_inplace_method(method_name):
                return self

            # Put back the wrappers where needed
            response = hook_args.hook_response(
                method_name,
                response,
                wrap_type=type(self),
                new_self=self,
                wrap_args=self.get_class_attributes(),
            )

        return response

    return overloaded_native_method