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

{leap_model_parser-0.1.250.dev2 → 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.dev2
+Version: 0.1.250.dev3
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT

{leap_model_parser-0.1.250.dev2 → 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
@@ -61,10 +62,11 @@ class ActivationInterpreter(fx.Interpreter):
     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()
+        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)
@@ -73,17 +75,16 @@ class ActivationInterpreter(fx.Interpreter):
 
 
 def get_module_activations(
-    ep: "ExportedProgram",
+    ep: ExportedProgram,
     user_inputs: List[torch.Tensor],
-) -> Dict[str, "np.ndarray"]:
-    import numpy as np
+) -> 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"] = {}
+    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()

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

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