<commit_before>import inspect
import re
import random
import logging
import types
import copy
import torch
import torch.nn as nn
from functools import wraps


import syft
from syft import workers
from syft.workers import BaseWorker
from .tensors.interpreters import TorchTensor
from .tensors.interpreters import PointerTensor
from .tensors.decorators import LoggingTensor
from .tensors.interpreters import FixedPrecisionTensor
from .tensors.interpreters import AdditiveSharingTensor
from .tensors.interpreters import MultiPointerTensor
from .torch_attributes import TorchAttributes
from .tensors.interpreters.abstract import initialize_tensor, _apply_args

from syft.exceptions import route_method_exception
from syft.exceptions import TensorsNotCollocatedException


class TorchHook:
    """A Hook which Overrides Methods on PyTorch Tensors.

    The purpose of this class is to:
        * extend torch methods to allow for the moving of tensors from one
        worker to another.
        * override torch methods to execute commands on one worker that are
        called on tensors controlled by the local worker.

    This class is typically the first thing you will initialize when using
    PySyft with PyTorch because it is responsible for augmenting PyTorch with
    PySyft's added functionality (such as remote execution).

    Args:
        local_worker: An optional BaseWorker instance that lets you provide a
            local worker as a parameter which TorchHook will assume to be the
            worker owned by the local machine. If you leave it empty,
            TorchClient will automatically initialize a
            :class:`.workers.VirtualWorker` under the assumption you're looking
            to do local experimentation or development.
        is_client: An optional boolean parameter (default True), indicating
            whether TorchHook is being initialized as an end-user client.This
            can impact whether or not variables are deleted when they fall out
            of scope. If you set this incorrectly on a end user client, Tensors
            and Variables will never be deleted. If you set this incorrectly on
            a remote machine (not a client), tensors will not get saved. It's
            really only important if you're not initializing the local worker
            yourself.
        verbose: An optional boolean parameter (default True) to indicate
            whether or not to print the operations as they occur.
        queue_size: An integer optional parameter (default 0) to specify the
            max length of the list that stores the messages to be sent.

    Example:
        >>> import syft as sy
        >>> hook = sy.TorchHook()
        Hooking into Torch...
        Overloading Complete.
        >>> x = sy.Tensor([-2,-1,0,1,2,3])
        >>> x
        -2
        -1
        0
        1
        2
        3
        [syft.core.frameworks.torch.tensor.FloatTensor of size 6]
    """

    def __init__(
        self, torch, local_worker: BaseWorker = None, is_client: bool = True, verbose: bool = True
    ):
        """Initializes the hook.

        Initialize the hook and define all the attributes pertaining to the
        torch hook in a special TorchAttibute class, that will be added in the
        syft.torch attributes. Hence, this parameters are now conveyed by the
        syft module.
        """
        # Save the provided torch module as an attribute of the hook
        self.torch = torch

        # Save the local worker as an attribute
        self.local_worker = local_worker

        if hasattr(torch, "torch_hooked"):
            logging.warning("Torch was already hooked... skipping hooking process")
            self.local_worker = syft.local_worker
            return
        else:
            torch.torch_hooked = True

        # Add all the torch attributes in the syft.torch attr
        syft.torch = TorchAttributes(torch, self)

        if self.local_worker is None:
            # Every TorchHook instance should have a local worker which is
            # responsible for interfacing with other workers. The worker
            # interface is what allows the Torch specific code in TorchHook to
            # be agnostic to the means by which workers communicate (such as
            # peer-to-peer, sockets, through local ports, or all within the
            # same process)
            self.local_worker = workers.VirtualWorker(
                hook=self, is_client_worker=is_client, id="me"
            )
        else:
            self.local_worker.hook = self

        self.to_auto_overload = {}

        self.args_hook_for_overloaded_attr = {}

        self._hook_native_tensor(torch.Tensor, TorchTensor)

        # Add all hooked tensor methods to pointer but change behaviour to have the cmd sent
        self._hook_pointer_tensor_methods()

        # Add all hooked tensor methods to multi_pointer to change behavior to hav ethe cmd
        # sent to all child pointers.
        self._hook_multi_pointer_tensor_methods()

        # Add all hooked tensor methods to Logging tensor but change behaviour to just forward
        # the cmd to the next child (behaviour can be changed in the SyftTensor class file)
        self._hook_syft_tensor_methods(LoggingTensor)

        # Add all hooked tensor methods to FixedPrecisionTensor tensor but change behaviour
        # to just forward the cmd to the next child (behaviour can be changed in the
        # SyftTensor class file)
        self._hook_syft_tensor_methods(FixedPrecisionTensor)

        # Add all hooked tensor methods to AdditiveSharingTensor tensor but change behaviour
        # to just forward the cmd to the next child (behaviour can be changed in the
        # SyftTensor class file)
        self._hook_syft_tensor_methods(AdditiveSharingTensor)

        # Hook the tensor constructor function
        self._hook_tensor()

        # Hook the Parameter methods to store tensor chains in parameters
        self._hook_parameters()

        # Hook torch functions from modules like torch.add OR torch.nn.functional (containing relu, etc.)
        self._hook_torch_module()

        # Hook torch.nn (containing Linear and Convolution layers)
        self._hook_module()

        # Add the local_worker to syft so that it can be found if the hook is
        # called several times
        syft.local_worker = self.local_worker
        syft.hook = self

    def _hook_native_tensor(self, tensor_type: type, syft_type: type):
        """Adds PySyft Tensor Functionality to the given native tensor type.

        Overloads the given native Torch tensor to add PySyft Tensor
        Functionality. Overloading involves modifying the tensor type with
        PySyft's added functionality. You may read about what kind of
        modifications are made in the methods that this method calls.

        Args:
            tensor_type: The type of tensor being hooked (in this refactor
                this is only ever torch.Tensor, but in previous versions of
                PySyft this iterated over all tensor types.
            syft_type: The abstract type whose methods should all be added to
                the tensor_type class. In practice this is always TorchTensor.
                Read more about it there.
        """
        # Reinitialize init method of Torch tensor with Syft init
        self._add_registration_to___init__(tensor_type, torch_tensor=True)

        # Overload Torch tensor properties with Syft properties
        self._hook_properties(tensor_type)

        # Returns a list of methods to be overloaded, stored in the dict to_auto_overload
        # with tensor_type as a key
        self.to_auto_overload[tensor_type] = self._which_methods_should_we_auto_overload(
            tensor_type
        )

        # [We don't rename native methods as torch tensors are not hooked] Rename native functions
        # #self._rename_native_functions(tensor_type)

        # Overload auto overloaded with Torch methods
        self._add_methods_from__torch_tensor(tensor_type, syft_type)

        self._hook_native_methods(tensor_type)

    def _hook_native_methods(self, tensor_type: type):
        """
        Add hooked version of all methods of the tensor_type to the
        Pointer tensor: instead of performing the native tensor
        method, it will be sent remotely to the location the pointer
        is pointing at.
        :param tensor_type: the tensor_type which holds the methods
        """
        # # Add methods defined in the TorchTensor class to the Pointer class
        # self._add_methods_from__torch_tensor(PointerTensor, TorchTensor)

        # Use a pre-defined list to select the methods to overload
        for attr in self.to_auto_overload[tensor_type]:
            # if we haven't already overloaded this function
            if f"native_{attr}" not in dir(tensor_type):
                native_method = getattr(tensor_type, attr)
                setattr(tensor_type, f"native_{attr}", native_method)
                new_method = self.get_hooked_method(attr)
                setattr(tensor_type, attr, new_method)

    def _hook_syft_tensor_methods(self, syft_type: type):
        """
        Add hooked version of all methods of the tensor_type to the
        Pointer tensor: instead of performing the native tensor
        method, it will be sent remotely to the location the pointer
        is pointing at.
        :param tensor_type: the tensor_type which holds the methods
        """

        tensor_type = self.torch.Tensor
        # Use a pre-defined list to select the methods to overload
        for attr in self.to_auto_overload[tensor_type]:
            if attr not in dir(syft_type):  # TODO and add special functions to include / avoid
                new_method = self.get_hooked_syft_method(attr)
                setattr(syft_type, attr, new_method)

    def _hook_pointer_tensor_methods(self):
        """
        Add hooked version of all methods of the tensor_type to the
        Pointer tensor: instead of performing the native tensor
        method, it will be sent remotely to the location the pointer
        is pointing at.
        """

        tensor_type = self.torch.Tensor
        # Use a pre-defined list to select the methods to overload
        for attr in self.to_auto_overload[tensor_type]:
            if attr not in dir(PointerTensor):  # TODO and add special functions to include / avoid
                new_method = self.get_hooked_pointer_method(attr)
                setattr(PointerTensor, attr, new_method)

    def _hook_multi_pointer_tensor_methods(self):
        """
        Add hooked version of all methods of the tensor_type to the
        Pointer tensor: instead of performing the native tensor
        method, it will be sent remotely to the location the pointer
        is pointing at.
        """

        tensor_type = self.torch.Tensor
        # Use a pre-defined list to select the methods to overload
        for attr in self.to_auto_overload[tensor_type]:
            if attr not in dir(
                MultiPointerTensor
            ):  # TODO and add special functions to include / avoid
                new_method = self.get_hooked_multi_pointer_method(attr)
                setattr(MultiPointerTensor, attr, new_method)

    def _hook_parameters(self):
        """
        This method overrides the torch Parameter class such that
        it works correctly with our overridden tensor types. The
        native torch Parameter class kept deleting all of our
        attributes on our custom tensors, so we wrote our own.
        """

        # Hook __new__ to handle when non-pure torch tensors are given as data attribute

        def hooked__new__(cls, data=None, requires_grad=True):
            if data is None:
                data = torch.Tensor()

            # If data is not a pure torch tensor you need to store the chain in a
            # specific place otherwise it will get deleted
            if not isinstance(data, torch.Tensor) or hasattr(data, "child"):
                p = torch.Tensor._make_subclass(cls, torch.Tensor(), requires_grad)
                if isinstance(data, torch.Tensor):  # so it's a wrapper: remove it
                    p.child = data.child
                else:
                    p.child = data
            else:
                p = torch.Tensor._make_subclass(cls, data, requires_grad)

            return p

        torch.nn.Parameter.__new__ = hooked__new__

        # Hook __repr__ to handle chain repr when needed

        torch.nn.Parameter.native_param___repr__ = torch.nn.Parameter.__repr__

        def hooked__repr__(self):
            if hasattr(self, "child"):
                return "Parameter containing:\n" + self.child.__repr__()
            else:
                return self.native_param___repr__()

        # torch.nn.Parameter.__repr__ = hooked__repr__

        # Hook .data to handle chain assignment when needed

        torch.nn.Parameter.native_param_data = torch.nn.Parameter.data

        @property
        def data(self):

            if hasattr(self, "child"):
                to_return = self.child.attr("data")
            else:

                to_return = self.native_param_data

                # good to ensure that the ID stays consistent
                # not 100% this is required but it's at least
                # good practice
                try:
                    to_return.id = self.data_id
                except:
                    self.data_id = to_return.id

            return to_return

        @data.setter
        def data(self, new_data):

            # If data is not a pure torch tensor you need to store the chain in a
            # specific place otherwise it will get deleted
            if not isinstance(new_data, torch.Tensor) or hasattr(new_data, "child"):
                self.child = new_data  # .wrap()
            else:
                if hasattr(self, "child"):
                    del self.child

                self.native_param_data.set_(new_data)  # .wrap()
            return self

        torch.nn.Parameter.data = data

        # Hook .grad to handle chain assignment when needed

        torch.nn.Parameter.native_param_grad = torch.nn.Parameter.grad

        @property
        def grad(self):

            if hasattr(self, "child"):
                to_return = self.child.attr("grad")
                if isinstance(to_return.child, syft.PointerTensor):
                    if to_return.child.is_none():
                        to_return = None
            else:
                to_return = self.native_param_grad

                # good to ensure that the ID stays consistent
                # not 100% this is required but it's at least
                # good practice
                try:
                    to_return.id = self.grad_id
                except AttributeError:
                    if to_return is not None and hasattr(to_return, "id"):
                        self.grad_id = to_return.id

            return to_return

        @grad.setter
        def grad(self, new_grad):

            # If grad is not a pure torch tensor you need to store the chain in a
            # specific place otherwise it will get deleted
            if new_grad is not None and (
                not isinstance(new_grad, torch.Tensor) or hasattr(new_grad, "child")
            ):
                self.child.grad = new_grad  # .wrap()
            else:
                if self.native_param_grad is not None:
                    self.native_param_grad.set_(new_grad)  # .wrap()
                elif new_grad is not None:
                    self.native_param_grad = new_grad
            return self

        torch.nn.Parameter.grad = grad

    def _hook_torch_module(self):
        """Overloads functions in the main torch modules.
        The way this is accomplished is by first moving all existing module
        functions in the torch module to native_<function_name_here>.

        Example:
            the real :func:`torch.cat` will become :func:`torch.native_cat`
            and :func:`torch.cat` will have our hooking code.
        """

        def perform_overloading(torch_module, func):

            # Where the overloading happens
            # 1. Get native function
            native_func = getattr(torch_module, func)
            # 2. Check it is a proper function
            if type(native_func) in [types.FunctionType, types.BuiltinFunctionType]:
                # 3. Build the hooked function
                new_func = self.get_hooked_func(native_func)
                # 4. Move the native function to its original module
                # /!\ Can be different from the torch_module!
                # Ex: in torch.py `torch.argmax = torch.functional.argmax`
                # ... So torch.argmax.__module__ is 'torch.functional' != 'torch'
                setattr(eval(native_func.__module__), f"native_{func}", native_func)
                # 5. Put instead the hooked one
                setattr(torch_module, func, new_func)

        # Hard fix for PyTorch versions < 1.0.2
        syft.torch.apply_fix16922(self.torch)
<commit_msg>Revert fix on module name for hooked functionc<commit_after>import inspect
import re
import random
import logging
import types
import copy
import torch
import torch.nn as nn
from functools import wraps


import syft
from syft import workers
from syft.workers import BaseWorker
from .tensors.interpreters import TorchTensor
from .tensors.interpreters import PointerTensor
from .tensors.decorators import LoggingTensor
from .tensors.interpreters import FixedPrecisionTensor
from .tensors.interpreters import AdditiveSharingTensor
from .tensors.interpreters import MultiPointerTensor
from .torch_attributes import TorchAttributes
from .tensors.interpreters.abstract import initialize_tensor, _apply_args

from syft.exceptions import route_method_exception
from syft.exceptions import TensorsNotCollocatedException


class TorchHook:
    """A Hook which Overrides Methods on PyTorch Tensors.

    The purpose of this class is to:
        * extend torch methods to allow for the moving of tensors from one
        worker to another.
        * override torch methods to execute commands on one worker that are
        called on tensors controlled by the local worker.

    This class is typically the first thing you will initialize when using
    PySyft with PyTorch because it is responsible for augmenting PyTorch with
    PySyft's added functionality (such as remote execution).

    Args:
        local_worker: An optional BaseWorker instance that lets you provide a
            local worker as a parameter which TorchHook will assume to be the
            worker owned by the local machine. If you leave it empty,
            TorchClient will automatically initialize a
            :class:`.workers.VirtualWorker` under the assumption you're looking
            to do local experimentation or development.
        is_client: An optional boolean parameter (default True), indicating
            whether TorchHook is being initialized as an end-user client.This
            can impact whether or not variables are deleted when they fall out
            of scope. If you set this incorrectly on a end user client, Tensors
            and Variables will never be deleted. If you set this incorrectly on
            a remote machine (not a client), tensors will not get saved. It's
            really only important if you're not initializing the local worker
            yourself.
        verbose: An optional boolean parameter (default True) to indicate
            whether or not to print the operations as they occur.
        queue_size: An integer optional parameter (default 0) to specify the
            max length of the list that stores the messages to be sent.

    Example:
        >>> import syft as sy
        >>> hook = sy.TorchHook()
        Hooking into Torch...
        Overloading Complete.
        >>> x = sy.Tensor([-2,-1,0,1,2,3])
        >>> x
        -2
        -1
        0
        1
        2
        3
        [syft.core.frameworks.torch.tensor.FloatTensor of size 6]
    """

    def __init__(
        self, torch, local_worker: BaseWorker = None, is_client: bool = True, verbose: bool = True
    ):
        """Initializes the hook.

        Initialize the hook and define all the attributes pertaining to the
        torch hook in a special TorchAttibute class, that will be added in the
        syft.torch attributes. Hence, this parameters are now conveyed by the
        syft module.
        """
        # Save the provided torch module as an attribute of the hook
        self.torch = torch

        # Save the local worker as an attribute
        self.local_worker = local_worker

        if hasattr(torch, "torch_hooked"):
            logging.warning("Torch was already hooked... skipping hooking process")
            self.local_worker = syft.local_worker
            return
        else:
            torch.torch_hooked = True

        # Add all the torch attributes in the syft.torch attr
        syft.torch = TorchAttributes(torch, self)

        if self.local_worker is None:
            # Every TorchHook instance should have a local worker which is
            # responsible for interfacing with other workers. The worker
            # interface is what allows the Torch specific code in TorchHook to
            # be agnostic to the means by which workers communicate (such as
            # peer-to-peer, sockets, through local ports, or all within the
            # same process)
            self.local_worker = workers.VirtualWorker(
                hook=self, is_client_worker=is_client, id="me"
            )
        else:
            self.local_worker.hook = self

        self.to_auto_overload = {}

        self.args_hook_for_overloaded_attr = {}

        self._hook_native_tensor(torch.Tensor, TorchTensor)

        # Add all hooked tensor methods to pointer but change behaviour to have the cmd sent
        self._hook_pointer_tensor_methods()

        # Add all hooked tensor methods to multi_pointer to change behavior to hav ethe cmd
        # sent to all child pointers.
        self._hook_multi_pointer_tensor_methods()

        # Add all hooked tensor methods to Logging tensor but change behaviour to just forward
        # the cmd to the next child (behaviour can be changed in the SyftTensor class file)
        self._hook_syft_tensor_methods(LoggingTensor)

        # Add all hooked tensor methods to FixedPrecisionTensor tensor but change behaviour
        # to just forward the cmd to the next child (behaviour can be changed in the
        # SyftTensor class file)
        self._hook_syft_tensor_methods(FixedPrecisionTensor)

        # Add all hooked tensor methods to AdditiveSharingTensor tensor but change behaviour
        # to just forward the cmd to the next child (behaviour can be changed in the
        # SyftTensor class file)
        self._hook_syft_tensor_methods(AdditiveSharingTensor)

        # Hook the tensor constructor function
        self._hook_tensor()

        # Hook the Parameter methods to store tensor chains in parameters
        self._hook_parameters()

        # Hook torch functions from modules like torch.add OR torch.nn.functional (containing relu, etc.)
        self._hook_torch_module()

        # Hook torch.nn (containing Linear and Convolution layers)
        self._hook_module()

        # Add the local_worker to syft so that it can be found if the hook is
        # called several times
        syft.local_worker = self.local_worker
        syft.hook = self

    def _hook_native_tensor(self, tensor_type: type, syft_type: type):
        """Adds PySyft Tensor Functionality to the given native tensor type.

        Overloads the given native Torch tensor to add PySyft Tensor
        Functionality. Overloading involves modifying the tensor type with
        PySyft's added functionality. You may read about what kind of
        modifications are made in the methods that this method calls.

        Args:
            tensor_type: The type of tensor being hooked (in this refactor
                this is only ever torch.Tensor, but in previous versions of
                PySyft this iterated over all tensor types.
            syft_type: The abstract type whose methods should all be added to
                the tensor_type class. In practice this is always TorchTensor.
                Read more about it there.
        """
        # Reinitialize init method of Torch tensor with Syft init
        self._add_registration_to___init__(tensor_type, torch_tensor=True)

        # Overload Torch tensor properties with Syft properties
        self._hook_properties(tensor_type)

        # Returns a list of methods to be overloaded, stored in the dict to_auto_overload
        # with tensor_type as a key
        self.to_auto_overload[tensor_type] = self._which_methods_should_we_auto_overload(
            tensor_type
        )

        # [We don't rename native methods as torch tensors are not hooked] Rename native functions
        # #self._rename_native_functions(tensor_type)

        # Overload auto overloaded with Torch methods
        self._add_methods_from__torch_tensor(tensor_type, syft_type)

        self._hook_native_methods(tensor_type)

    def _hook_native_methods(self, tensor_type: type):
        """
        Add hooked version of all methods of the tensor_type to the
        Pointer tensor: instead of performing the native tensor
        method, it will be sent remotely to the location the pointer
        is pointing at.
        :param tensor_type: the tensor_type which holds the methods
        """
        # # Add methods defined in the TorchTensor class to the Pointer class
        # self._add_methods_from__torch_tensor(PointerTensor, TorchTensor)

        # Use a pre-defined list to select the methods to overload
        for attr in self.to_auto_overload[tensor_type]:
            # if we haven't already overloaded this function
            if f"native_{attr}" not in dir(tensor_type):
                native_method = getattr(tensor_type, attr)
                setattr(tensor_type, f"native_{attr}", native_method)
                new_method = self.get_hooked_method(attr)
                setattr(tensor_type, attr, new_method)

    def _hook_syft_tensor_methods(self, syft_type: type):
        """
        Add hooked version of all methods of the tensor_type to the
        Pointer tensor: instead of performing the native tensor
        method, it will be sent remotely to the location the pointer
        is pointing at.
        :param tensor_type: the tensor_type which holds the methods
        """

        tensor_type = self.torch.Tensor
        # Use a pre-defined list to select the methods to overload
        for attr in self.to_auto_overload[tensor_type]:
            if attr not in dir(syft_type):  # TODO and add special functions to include / avoid
                new_method = self.get_hooked_syft_method(attr)
                setattr(syft_type, attr, new_method)

    def _hook_pointer_tensor_methods(self):
        """
        Add hooked version of all methods of the tensor_type to the
        Pointer tensor: instead of performing the native tensor
        method, it will be sent remotely to the location the pointer
        is pointing at.
        """

        tensor_type = self.torch.Tensor
        # Use a pre-defined list to select the methods to overload
        for attr in self.to_auto_overload[tensor_type]:
            if attr not in dir(PointerTensor):  # TODO and add special functions to include / avoid
                new_method = self.get_hooked_pointer_method(attr)
                setattr(PointerTensor, attr, new_method)

    def _hook_multi_pointer_tensor_methods(self):
        """
        Add hooked version of all methods of the tensor_type to the
        Pointer tensor: instead of performing the native tensor
        method, it will be sent remotely to the location the pointer
        is pointing at.
        """

        tensor_type = self.torch.Tensor
        # Use a pre-defined list to select the methods to overload
        for attr in self.to_auto_overload[tensor_type]:
            if attr not in dir(
                MultiPointerTensor
            ):  # TODO and add special functions to include / avoid
                new_method = self.get_hooked_multi_pointer_method(attr)
                setattr(MultiPointerTensor, attr, new_method)

    def _hook_parameters(self):
        """
        This method overrides the torch Parameter class such that
        it works correctly with our overridden tensor types. The
        native torch Parameter class kept deleting all of our
        attributes on our custom tensors, so we wrote our own.
        """

        # Hook __new__ to handle when non-pure torch tensors are given as data attribute

        def hooked__new__(cls, data=None, requires_grad=True):
            if data is None:
                data = torch.Tensor()

            # If data is not a pure torch tensor you need to store the chain in a
            # specific place otherwise it will get deleted
            if not isinstance(data, torch.Tensor) or hasattr(data, "child"):
                p = torch.Tensor._make_subclass(cls, torch.Tensor(), requires_grad)
                if isinstance(data, torch.Tensor):  # so it's a wrapper: remove it
                    p.child = data.child
                else:
                    p.child = data
            else:
                p = torch.Tensor._make_subclass(cls, data, requires_grad)

            return p

        torch.nn.Parameter.__new__ = hooked__new__

        # Hook __repr__ to handle chain repr when needed

        torch.nn.Parameter.native_param___repr__ = torch.nn.Parameter.__repr__

        def hooked__repr__(self):
            if hasattr(self, "child"):
                return "Parameter containing:\n" + self.child.__repr__()
            else:
                return self.native_param___repr__()

        # torch.nn.Parameter.__repr__ = hooked__repr__

        # Hook .data to handle chain assignment when needed

        torch.nn.Parameter.native_param_data = torch.nn.Parameter.data

        @property
        def data(self):

            if hasattr(self, "child"):
                to_return = self.child.attr("data")
            else:

                to_return = self.native_param_data

                # good to ensure that the ID stays consistent
                # not 100% this is required but it's at least
                # good practice
                try:
                    to_return.id = self.data_id
                except:
                    self.data_id = to_return.id

            return to_return

        @data.setter
        def data(self, new_data):

            # If data is not a pure torch tensor you need to store the chain in a
            # specific place otherwise it will get deleted
            if not isinstance(new_data, torch.Tensor) or hasattr(new_data, "child"):
                self.child = new_data  # .wrap()
            else:
                if hasattr(self, "child"):
                    del self.child

                self.native_param_data.set_(new_data)  # .wrap()
            return self

        torch.nn.Parameter.data = data

        # Hook .grad to handle chain assignment when needed

        torch.nn.Parameter.native_param_grad = torch.nn.Parameter.grad

        @property
        def grad(self):

            if hasattr(self, "child"):
                to_return = self.child.attr("grad")
                if isinstance(to_return.child, syft.PointerTensor):
                    if to_return.child.is_none():
                        to_return = None
            else:
                to_return = self.native_param_grad

                # good to ensure that the ID stays consistent
                # not 100% this is required but it's at least
                # good practice
                try:
                    to_return.id = self.grad_id
                except AttributeError:
                    if to_return is not None and hasattr(to_return, "id"):
                        self.grad_id = to_return.id

            return to_return

        @grad.setter
        def grad(self, new_grad):

            # If grad is not a pure torch tensor you need to store the chain in a
            # specific place otherwise it will get deleted
            if new_grad is not None and (
                not isinstance(new_grad, torch.Tensor) or hasattr(new_grad, "child")
            ):
                self.child.grad = new_grad  # .wrap()
            else:
                if self.native_param_grad is not None:
                    self.native_param_grad.set_(new_grad)  # .wrap()
                elif new_grad is not None:
                    self.native_param_grad = new_grad
            return self

        torch.nn.Parameter.grad = grad

    def _hook_torch_module(self):
        """Overloads functions in the main torch modules.
        The way this is accomplished is by first moving all existing module
        functions in the torch module to native_<function_name_here>.

        Example:
            the real :func:`torch.cat` will become :func:`torch.native_cat`
            and :func:`torch.cat` will have our hooking code.
        """

        def perform_overloading(torch_module, func):

            # Where the overloading happens
            # 1. Get native function
            native_func = getattr(torch_module, func)
            # 2. Check it is a proper function
            if type(native_func) in [types.FunctionType, types.BuiltinFunctionType]:
                # 3. Build the hooked function
                new_func = self.get_hooked_func(native_func)
                # 4. Move the native function
                setattr(torch_module, f"native_{func}", native_func)
                # 5. Put instead the hooked one
                setattr(torch_module, func, new_func)

        # Hard fix for PyTorch versions < 1.0.2
        syft.torch.apply_fix16922(self.torch)

{'fraction_non_alphanumeric': 0.05373660255469094, 'fraction_numerical': 0.0019380414036118045, 'mean_word_length': 2.1312982714233173, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1183722', 'n_tokens_mistral': 8770, 'n_tokens_neox': 8401, 'n_words': 3329}