leap-model-parser 0.1.250.dev0__tar.gz → 0.1.250.dev2__tar.gz

{leap_model_parser-0.1.250.dev0 → leap_model_parser-0.1.250.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.250.dev0
+Version: 0.1.250.dev2
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT
@@ -17,7 +17,7 @@ Requires-Dist: keras-data-format-converter (==0.1.24)
 Requires-Dist: leap-model-rebuilder (==0.1.7)
 Requires-Dist: numpy (>=1.22.3,<2.0.0)
 Requires-Dist: onnx (==1.13.0)
-Requires-Dist: onnx2kerastl (==0.0.190)
+Requires-Dist: onnx2kerastl (==0.0.191)
 Requires-Dist: tensorflow (==2.12.0) ; platform_machine == "x86_64"
 Requires-Dist: tensorflow-io-gcs-filesystem (==0.34.0)
 Requires-Dist: tensorflow-macos (==2.12.0) ; platform_machine == "arm64"

{leap_model_parser-0.1.250.dev0 → leap_model_parser-0.1.250.dev2}/leap_model_parser/torch_utils.py

@@ -58,8 +58,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)) and len(value) > 0:
+            first = value[0]
+            if isinstance(first, torch.Tensor):
+                self.cache[name] = first.detach().clone()
+
     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"]:
+    import numpy as np
+    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.dev0 → leap_model_parser-0.1.250.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.250.dev"
+version = "0.1.250.dev2"
 description = ""
 authors = ["idan <idan.yogev@tensorleap.ai>"]
 license = "MIT"
@@ -18,7 +18,7 @@ tensorflow = {version = "2.12.0", markers = "platform_machine  == 'x86_64'"}
 tensorflow-macos = {version = "2.12.0", markers = "platform_machine  == 'arm64'"}
 numpy = "^1.22.3"
 onnx = "1.13.0"
-onnx2kerastl = "0.0.190"
+onnx2kerastl = "0.0.191"
 keras-data-format-converter = "0.1.24"
 leap-model-rebuilder = "0.1.7"
 tensorflow-io-gcs-filesystem = "0.34.0"