<commit_before>import torch
import syft as sy
from syft.exceptions import RemoteObjectFoundError
from syft.exceptions import PureTorchTensorFoundError

from syft.exceptions import ResponseSignatureError
from syft.frameworks.torch.tensors.interpreters import AutogradTensor
from syft.frameworks.torch.tensors.interpreters import AbstractTensor
from syft.frameworks.torch.pointers import PointerTensor
from syft.frameworks.torch.tensors.interpreters import TorchTensor
from syft.frameworks.torch.tensors.interpreters import FixedPrecisionTensor
from syft.frameworks.torch.tensors.interpreters import AdditiveSharingTensor
from syft.frameworks.torch.tensors.interpreters import MultiPointerTensor
from syft.frameworks.torch.tensors.interpreters import CRTPrecisionTensor
from syft.frameworks.torch.tensors.interpreters import LargePrecisionTensor
from syft.frameworks.torch.tensors.decorators import LoggingTensor

from typing import Callable
from typing import Union
from typing import Tuple
from typing import List

import numpy as np


hook_method_args_functions = {}
hook_method_response_functions = {}
get_tensor_type_functions = {}

base_types = {int, float, str, bool, bytes, bytearray, complex}

one = lambda _args: 1

# dict to specify the action depending of the type found
type_rule = {
    list: lambda _args: [build_rule(a) for a in _args],
    tuple: lambda _args: tuple([build_rule(a) for a in _args]),
    dict: one,  # FIXME This is for additiveShareTensor.child, it can be confusing and AST.child
    np.ndarray: one,
    # should perhaps be of type ShareDict extending dict or something like this
    LoggingTensor: one,
    FixedPrecisionTensor: one,
    AutogradTensor: one,
    AdditiveSharingTensor: one,
    MultiPointerTensor: one,
    CRTPrecisionTensor: one,
    PointerTensor: one,
    LargePrecisionTensor: one,
    torch.Tensor: one,
    torch.nn.Parameter: one,
}

# Dict to return the proper lambda function for the right torch or syft tensor type
forward_func = {
    PointerTensor: lambda p: (_ for _ in ()).throw(RemoteObjectFoundError(p)),
    torch.Tensor: lambda i: i.child
    if hasattr(i, "child")
    else (_ for _ in ()).throw(PureTorchTensorFoundError),
    torch.nn.Parameter: lambda i: i.child
    if hasattr(i, "child")
    else (_ for _ in ()).throw(PureTorchTensorFoundError),
    LoggingTensor: lambda i: i.child,
    FixedPrecisionTensor: lambda i: i.child,
    AutogradTensor: lambda i: i.child,
    AdditiveSharingTensor: lambda i: i.child,
    MultiPointerTensor: lambda i: i.child,
    CRTPrecisionTensor: lambda i: i.child,
    LargePrecisionTensor: lambda i: i._internal_representation_to_large_ints(),
    "my_syft_tensor_type": lambda i: i.child,
}

# Dict to return the proper lambda function for the right torch or syft tensor type
backward_func = {
    TorchTensor: lambda i: i.wrap(),
    torch.Tensor: lambda i: i.wrap(),
    torch.nn.Parameter: lambda i: torch.nn.Parameter(data=i),
    PointerTensor: lambda i: i,
    LoggingTensor: lambda i: LoggingTensor().on(i, wrap=False),
    FixedPrecisionTensor: lambda i, **kwargs: FixedPrecisionTensor(**kwargs).on(i, wrap=False),
    LargePrecisionTensor: lambda i, **kwargs: LargePrecisionTensor(**kwargs).on(
        LargePrecisionTensor.create_tensor_from_numpy(i, **kwargs), wrap=False
    ),
    AutogradTensor: lambda i: AutogradTensor(data=i).on(i, wrap=False),
    AdditiveSharingTensor: lambda i, **kwargs: AdditiveSharingTensor(**kwargs).on(i, wrap=False),
    MultiPointerTensor: lambda i, **kwargs: MultiPointerTensor(**kwargs).on(i, wrap=False),
    CRTPrecisionTensor: lambda i, **kwargs: CRTPrecisionTensor(**kwargs).on(i, wrap=False),
    "my_syft_tensor_type": lambda i, **kwargs: "my_syft_tensor_type(**kwargs).on(i, wrap=False)",
}

# Methods or functions whose signature changes a lot and that we don't want to "cache", because
# they have an arbitrary number of tensors in args which can trigger unexpected behaviour
ambiguous_methods = {"__getitem__", "_getitem_public", "view", "permute", "add_", "sub_"}
ambiguous_functions = {"torch.unbind", "unbind", "torch.stack", "stack", "torch.mean", "torch.sum"}


def unwrap_args_from_method(attr, method_self, args, kwargs):
    """Method arguments are sometimes simple types (such as strings or ints) but
    sometimes they are custom Syft tensors such as wrappers (torch.Tensor) or LoggingTensor
    or some other tensor type. Complex types (which have a .child attribute) need to
    have arguments converted from the arg to arg.child so that the types match as the
    method is being called down the chain. To make this efficient, we cache which args
    need to be replaced with their children in a dictionary called
    hook_method_args_functions. However, sometimes a method (an attr) has multiple
    different argument signatures, such that sometimes arguments have .child objects
    and other times they don't (such as x.div(), which can accept either a tensor or a
    float as an argument). This invalidates the cache, so we need to have a try/except
    which refreshes the cache if the signature triggers an error.

    Args:<commit_msg>added split in ambiguous functions (#2519)<commit_after>import torch
import syft as sy
from syft.exceptions import RemoteObjectFoundError
from syft.exceptions import PureTorchTensorFoundError

from syft.exceptions import ResponseSignatureError
from syft.frameworks.torch.tensors.interpreters import AutogradTensor
from syft.frameworks.torch.tensors.interpreters import AbstractTensor
from syft.frameworks.torch.pointers import PointerTensor
from syft.frameworks.torch.tensors.interpreters import TorchTensor
from syft.frameworks.torch.tensors.interpreters import FixedPrecisionTensor
from syft.frameworks.torch.tensors.interpreters import AdditiveSharingTensor
from syft.frameworks.torch.tensors.interpreters import MultiPointerTensor
from syft.frameworks.torch.tensors.interpreters import CRTPrecisionTensor
from syft.frameworks.torch.tensors.interpreters import LargePrecisionTensor
from syft.frameworks.torch.tensors.decorators import LoggingTensor

from typing import Callable
from typing import Union
from typing import Tuple
from typing import List

import numpy as np


hook_method_args_functions = {}
hook_method_response_functions = {}
get_tensor_type_functions = {}

base_types = {int, float, str, bool, bytes, bytearray, complex}

one = lambda _args: 1

# dict to specify the action depending of the type found
type_rule = {
    list: lambda _args: [build_rule(a) for a in _args],
    tuple: lambda _args: tuple([build_rule(a) for a in _args]),
    dict: one,  # FIXME This is for additiveShareTensor.child, it can be confusing and AST.child
    np.ndarray: one,
    # should perhaps be of type ShareDict extending dict or something like this
    LoggingTensor: one,
    FixedPrecisionTensor: one,
    AutogradTensor: one,
    AdditiveSharingTensor: one,
    MultiPointerTensor: one,
    CRTPrecisionTensor: one,
    PointerTensor: one,
    LargePrecisionTensor: one,
    torch.Tensor: one,
    torch.nn.Parameter: one,
}

# Dict to return the proper lambda function for the right torch or syft tensor type
forward_func = {
    PointerTensor: lambda p: (_ for _ in ()).throw(RemoteObjectFoundError(p)),
    torch.Tensor: lambda i: i.child
    if hasattr(i, "child")
    else (_ for _ in ()).throw(PureTorchTensorFoundError),
    torch.nn.Parameter: lambda i: i.child
    if hasattr(i, "child")
    else (_ for _ in ()).throw(PureTorchTensorFoundError),
    LoggingTensor: lambda i: i.child,
    FixedPrecisionTensor: lambda i: i.child,
    AutogradTensor: lambda i: i.child,
    AdditiveSharingTensor: lambda i: i.child,
    MultiPointerTensor: lambda i: i.child,
    CRTPrecisionTensor: lambda i: i.child,
    LargePrecisionTensor: lambda i: i._internal_representation_to_large_ints(),
    "my_syft_tensor_type": lambda i: i.child,
}

# Dict to return the proper lambda function for the right torch or syft tensor type
backward_func = {
    TorchTensor: lambda i: i.wrap(),
    torch.Tensor: lambda i: i.wrap(),
    torch.nn.Parameter: lambda i: torch.nn.Parameter(data=i),
    PointerTensor: lambda i: i,
    LoggingTensor: lambda i: LoggingTensor().on(i, wrap=False),
    FixedPrecisionTensor: lambda i, **kwargs: FixedPrecisionTensor(**kwargs).on(i, wrap=False),
    LargePrecisionTensor: lambda i, **kwargs: LargePrecisionTensor(**kwargs).on(
        LargePrecisionTensor.create_tensor_from_numpy(i, **kwargs), wrap=False
    ),
    AutogradTensor: lambda i: AutogradTensor(data=i).on(i, wrap=False),
    AdditiveSharingTensor: lambda i, **kwargs: AdditiveSharingTensor(**kwargs).on(i, wrap=False),
    MultiPointerTensor: lambda i, **kwargs: MultiPointerTensor(**kwargs).on(i, wrap=False),
    CRTPrecisionTensor: lambda i, **kwargs: CRTPrecisionTensor(**kwargs).on(i, wrap=False),
    "my_syft_tensor_type": lambda i, **kwargs: "my_syft_tensor_type(**kwargs).on(i, wrap=False)",
}

# Methods or functions whose signature changes a lot and that we don't want to "cache", because
# they have an arbitrary number of tensors in args which can trigger unexpected behaviour
ambiguous_methods = {"__getitem__", "_getitem_public", "view", "permute", "add_", "sub_"}
ambiguous_functions = {
    "torch.unbind",
    "unbind",
    "torch.stack",
    "stack",
    "torch.mean",
    "torch.sum",
    "torch.functional.split",
    "split",
}


def unwrap_args_from_method(attr, method_self, args, kwargs):
    """Method arguments are sometimes simple types (such as strings or ints) but
    sometimes they are custom Syft tensors such as wrappers (torch.Tensor) or LoggingTensor
    or some other tensor type. Complex types (which have a .child attribute) need to
    have arguments converted from the arg to arg.child so that the types match as the
    method is being called down the chain. To make this efficient, we cache which args
    need to be replaced with their children in a dictionary called
    hook_method_args_functions. However, sometimes a method (an attr) has multiple
    different argument signatures, such that sometimes arguments have .child objects
    and other times they don't (such as x.div(), which can accept either a tensor or a
    float as an argument). This invalidates the cache, so we need to have a try/except
    which refreshes the cache if the signature triggers an error.

    Args:
{'fraction_non_alphanumeric': 0.08712450497440355, 'fraction_numerical': 0.0005795421616922632, 'mean_word_length': 4.726769911504425, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 3, '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': '1200248', 'n_tokens_mistral': 3081, 'n_tokens_neox': 3041, 'n_words': 1030}