Proposal: Registry Pattern for mapping between f and jet_f

[TimSiebert1]

We could use the following registry that stores our f->jet_f

@dataclass
class JetRegistry:
    _map: dict[Callable, Callable] = field(default_factory=dict)
    _frozen: bool = False

    def register(self, *funcs: Callable, allow_override: bool = False):
        def deco(jet_f: Callable):
            for f in funcs:
                if not allow_override and f in self._map:
                    raise RuntimeError(
                        f"Jet for {getattr(f,'__name__',f)} already exists"
                    )
                self._map[f] = jet_f
            return jet_f
        return deco

    def get(self, f: Callable) -> Callable:
        return self._map[f]

    def mapping(self) -> dict[Callable, Callable]:
        return dict(self._map)  # copy

""" 
maybe also the possibility to freeze the registry
    def freeze(self):
        self._frozen = True
"""

If JET = JetRegistry(), users could extend the registry using

@JET.register(torch.sin, math.sin)
def jet_sin(
    s: PrimalAndCoefficients, K: int, is_taylor: tuple[bool, ...]
) -> ValueAndCoefficients:
    pass

What do you think?

[f-dangel]

I think using torch.fx.Transformer would be a good approach to simplify the graph overloading function.
We can use this for overloading for jet, but also for overloading for traceable_vmap.

Example:

from torch.fx import Transformer

class ReluToSigmoid(Transformer):
    def call_function(self, target, args, kwargs):
        if target == F.relu:
            return super().call_function(torch.sigmoid, args, kwargs)
        return super().call_function(target, args, kwargs)

sigmoid_gm = ReluToSigmoid(gm).transform()

[f-dangel]

I think the fx.Transformer-based refactoring only makes sense after standardizing the signatures of jet_* operations (#91)

[TimSiebert1]

What about something like this:

from torch.fx import Transformer, GraphModule
from jet.operations import jet_sin, MAPPING, JetInfo


class SinToJet(Transformer):
    def __init__(self, module, derivative_order):
        super().__init__(module)
        self.derivative_order = derivative_order

    def call_function(self, target, args, kwargs):
        if target in MAPPING.keys():
            is_taylor = [... for arg in args]
            super().call_function(
                jet_sin,
                args,
                JetInfo(derivative_order=self.derivative_order, is_taylor=is_taylor),
            )
        return super().call_function(target, args, kwargs)


transformations = [SinToJet]

We could replace the logic in _replace_operations_with_taylor using something like this

for transformation in transformations:
    gm = transformation(gm, derivative_order).transform()