[Func] Allow proper preservations/reference across copy and pickle

[vasqu]

As per title, we want copy/deepcopy as well as pickle to keep proper references to the actual functions. This is currently failing / losing track of the correct references. The original issue was noticed by @SunMarc see repro

import copy
import pickle

from torch import nn

from kernels import Mode, kernelize, use_kernel_mapping
from kernels import use_kernel_forward_from_hub, use_kernelized_func


@use_kernel_forward_from_hub("f")
def f(x):
    return x + 1


@use_kernelized_func(f)
class Model(nn.Module):
    def forward(self, x):
        return f(x)


def add_ten(x):
    return x + 10


a, b = Model(), Model()

# 1. Simulate swapping A's function: B changes too.
a._kernel_funcs["f"].forward = add_ten
print(a(1), b(1))  # 11, 11

# 2. Reset a copy: its forward still calls the global wrapper.
c = copy.deepcopy(a)
with use_kernel_mapping({}, inherit_mapping=False):
    kernelize(c, device="cpu", mode=Mode.INFERENCE)

print(c._kernel_funcs["f"](1), c(1), a(1))  # 2, 11, 11

# 3. Pickle doesn't work
d = pickle.loads(pickle.dumps(a))
# AttributeError: Can't pickle local object '_create_func_module.<locals>.Func'

[vasqu]

cc @danieldk

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[vasqu]

Still draft, noticed we actually dont want singleton but exactly the reverse

[vasqu]

pickle is the complex case as this is just due to pickle 😅 I think we can keep it since some do use pickle but they should use safetensors

[vasqu]

it's needed for torch compile but i think there are works to make it native within dynamo, still need it for BC either way

[vasqu]

also let me know whether to move things to other files as im not super familiar with the repo structure I thought its simpler to just add it directly for now