leap-model-parser 0.1.250.dev1__tar.gz → 0.1.250.dev3__tar.gz

{leap_model_parser-0.1.250.dev1 → leap_model_parser-0.1.250.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.250.dev1
+Version: 0.1.250.dev3
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT

{leap_model_parser-0.1.250.dev1 → leap_model_parser-0.1.250.dev3}/leap_model_parser/torch_utils.py

@@ -1,3 +1,4 @@
+import numpy as np
 import torch
 import torch.fx as fx
 from torch.export import ExportedProgram
@@ -58,8 +59,33 @@ class ActivationInterpreter(fx.Interpreter):
         super().__init__(gm)
         self.cache: Dict[str, torch.Tensor] = {}
 
+    def _store_tensor(self, name: str, value: Any) -> None:
+        if isinstance(value, torch.Tensor):
+            self.cache[name] = value.detach().clone()
+        elif isinstance(value, (tuple, list)):
+            for item in value:
+                if isinstance(item, torch.Tensor):
+                    self.cache[name] = item.detach().clone()
+                    break
+
     def run_node(self, node: fx.Node) -> Any:
         result = super().run_node(node)
-        if isinstance(result, torch.Tensor):
-            self.cache[node.name] = result.detach().clone()
+        self._store_tensor(node.name, result)
         return result
+
+
+def get_module_activations(
+    ep: ExportedProgram,
+    user_inputs: List[torch.Tensor],
+) -> Dict[str, np.ndarray]:
+    tensors = [torch.from_numpy(x.astype(np.float32)) if isinstance(x, np.ndarray) else x
+               for x in user_inputs]
+    m2n = build_module_to_node_map(ep)
+    interp = ActivationInterpreter(ep.graph_module)
+    with torch.no_grad():
+        interp.run(*get_full_args(ep, tensors))
+    result: Dict[str, np.ndarray] = {}
+    for mod_path, (node_name, _) in m2n.items():
+        if node_name in interp.cache:
+            result[mod_path] = interp.cache[node_name].numpy()
+    return result

{leap_model_parser-0.1.250.dev1 → leap_model_parser-0.1.250.dev3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.250.dev1"
+version = "0.1.250.dev3"
 description = ""
 authors = ["idan <idan.yogev@tensorleap.ai>"]
 license = "MIT"