Merge pull request #285 from ndif-team/local

Local

src/nnsight/__init__.py

@@ -84,11 +84,14 @@ def inner(cls, fake_mode, elem, device, constant=None):
 
 DEFAULT_PATCHER.__enter__()
 
-from .tracing.contexts import GlobalTracingContext
-
-apply = GlobalTracingContext.GLOBAL_TRACING_CONTEXT.apply
-log = GlobalTracingContext.GLOBAL_TRACING_CONTEXT.log
+from .intervention.contexts import GlobalInterventionTracingContext
 
+apply = GlobalInterventionTracingContext.GLOBAL_TRACING_CONTEXT.apply
+log = GlobalInterventionTracingContext.GLOBAL_TRACING_CONTEXT.log
+local = GlobalInterventionTracingContext.GLOBAL_TRACING_CONTEXT.local
+cond = GlobalInterventionTracingContext.GLOBAL_TRACING_CONTEXT.cond
+iter = GlobalInterventionTracingContext.GLOBAL_TRACING_CONTEXT.iter
+stop = GlobalInterventionTracingContext.GLOBAL_TRACING_CONTEXT.stop
 
 def trace(fn):
     """Helper decorator to add a function to the intervention graph via `.apply(...)`.
@@ -103,7 +106,7 @@ def trace(fn):
 
     @wraps(fn)
     def inner(*args, **kwargs):
-
+        
         return apply(fn, *args, **kwargs)
 
     return inner
@@ -127,37 +130,6 @@ def inner(*args, **kwargs):
 from .intervention.graph import InterventionProxy
 
 
-# def local(object: Callable | InterventionProxy):
-#     """Helper decorator to add a function to the intervention graph via `.apply(...)`
-#     AND convert all input Proxies to local ones via `.local()`.
-
-#     If a non-function is passed in, its assumed to be an `InterventionProxy` and `.local()` is called and returned.
-
-#     Args:
-#         object ( Callable | InterventionProxy): Function to apply or Proxy to make local.
-
-#     Returns:
-#         Callable | InterventionProxy: Traceable local function or local Proxy.
-#     """
-
-#     if inspect.isroutine(object):
-
-#         fn = trace(object)
-
-#         @wraps(fn)
-#         def inner(*args, **kwargs):
-
-#             args, kwargs = util.apply(
-#                 (args, kwargs), lambda x: x.local(), InterventionProxy
-#             )
-
-#             return fn(*args, **kwargs)
-
-#         return inner
-
-#     return object.local()
-
-
 # def remote(object: Callable | Any):
 #     """Helper decorator to add a function to the intervention graph via `.apply(...)`
 #     AND convert all input Proxies to downloaded local ones via `.local()`

src/nnsight/config.yaml

@@ -1,10 +1,10 @@
 API:
-  APIKEY: null
+  APIKEY: rhJV6e3LGhtXlkB47z9W
+  FORMAT: json
   HOST: ndif.dev
   JOB_ID: null
   SSL: true
   ZLIB: true
-  FORMAT: json
 APP:
   LOGGING: false
   REMOTE_LOGGING: true

src/nnsight/intervention/__init__.py

@@ -1,2 +1,10 @@
+
+"""
+The `intervention` module extends the `tracing` module to add PyTorch specific interventions to a given computation graph.
+It defines its own: protocols, contexts, backends and graph primitives to achieve this.
+
+
+
+"""
 from .base import NNsight
 from .envoy import Envoy