Add helper function to replace whitelisted functions

src/nnsight/schema/format/functions.py

@@ -1,5 +1,6 @@
 import operator
 from inspect import getmembers, isbuiltin, isfunction, ismethod, ismethoddescriptor, isclass
+from typing import Callable
 
 import einops
 import torch
@@ -24,6 +25,15 @@ def get_function_name(fn, module_name=None):
 
     return f"{module_name}.{fn.__qualname__}"
 
+def update_function(function: str | Callable, new_function: Callable):
+
+    if not isinstance(function, str):
+
+        function = get_function_name(function)
+
+    new_function.__name__ = function
+
+    FUNCTIONS_WHITELIST[function] = new_function
 
 FUNCTIONS_WHITELIST = {}