leap-model-parser 0.1.250.dev5__tar.gz → 0.1.250.dev7__tar.gz

{leap_model_parser-0.1.250.dev5 → leap_model_parser-0.1.250.dev7}/PKG-INFO

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

{leap_model_parser-0.1.250.dev5 → leap_model_parser-0.1.250.dev7}/leap_model_parser/torch_utils.py

@@ -3,7 +3,7 @@ import torch
 import torch.fx as fx
 from torch.export import ExportedProgram
 from torch.export.graph_signature import InputKind
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 
 def get_full_args(ep: ExportedProgram, user_inputs: List[torch.Tensor]) -> List[torch.Tensor]:
@@ -55,17 +55,28 @@ def extract_tensor(output: Any) -> torch.Tensor:
 
 
 class ActivationInterpreter(fx.Interpreter):
-    def __init__(self, gm: fx.GraphModule):
+    def __init__(self, gm: fx.GraphModule, grad_mode: bool = False):
         super().__init__(gm)
         self.cache: Dict[str, torch.Tensor] = {}
+        self._grad_mode = grad_mode
 
     def _store_tensor(self, name: str, value: Any) -> None:
         if isinstance(value, torch.Tensor):
-            self.cache[name] = value.detach().clone()
+            if self._grad_mode:
+                self.cache[name] = value
+                if value.requires_grad:
+                    value.retain_grad()
+            else:
+                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()
+                    if self._grad_mode:
+                        self.cache[name] = item
+                        if item.requires_grad:
+                            item.retain_grad()
+                    else:
+                        self.cache[name] = item.detach().clone()
                     break
 
     def run_node(self, node: fx.Node) -> Any:
@@ -74,6 +85,50 @@ class ActivationInterpreter(fx.Interpreter):
         return result
 
 
+def get_pt2_output_module_paths(ep: ExportedProgram) -> List[Optional[str]]:
+    """Return the nn.Module path responsible for each output slot in the FX graph."""
+    output_node = None
+    for node in ep.graph_module.graph.nodes:
+        if node.op == "output":
+            output_node = node
+            break
+    if output_node is None:
+        return []
+
+    args = output_node.args[0]
+    if not isinstance(args, (list, tuple)):
+        args = [args]
+
+    result: List[Optional[str]] = []
+    for output_ref in args:
+        if output_ref is None or not hasattr(output_ref, "meta"):
+            result.append(None)
+            continue
+
+        visited: set = set()
+        current = output_ref
+        mod_path = None
+        while current.op == "call_function" and current not in visited:
+            stack = current.meta.get("nn_module_stack", {})
+            if stack:
+                mod_path = list(stack.items())[-1][1][0]
+                break
+            if current.args:
+                visited.add(current)
+                current = current.args[0]
+            else:
+                break
+
+        if mod_path is None:
+            stack = current.meta.get("nn_module_stack", {})
+            if stack:
+                mod_path = list(stack.items())[-1][1][0]
+
+        result.append(mod_path)
+
+    return result
+
+
 def get_module_activations(
     ep: ExportedProgram,
     user_inputs: List[torch.Tensor],

{leap_model_parser-0.1.250.dev5 → leap_model_parser-0.1.250.dev7}/pyproject.toml

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