leap-model-parser 0.1.250.dev4__tar.gz → 0.1.250.dev6__tar.gz

{leap_model_parser-0.1.250.dev4 → leap_model_parser-0.1.250.dev6}/PKG-INFO

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

{leap_model_parser-0.1.250.dev4 → leap_model_parser-0.1.250.dev6}/leap_model_parser/torch_graph_editor.py

@@ -28,6 +28,9 @@ class TorchLeapGraphEditor(LeapGraphEditor):
             )
         return module_path
 
+    def _handle_input_node_with_index(self, input_node) -> str:
+        return self._add_input_encoder_not_connected_to_the_model_node(input_node.name)
+
     def _resolve_input_origin_name(self, input_index: int) -> str:
         from torch.export.graph_signature import InputKind
         user_inputs = []

{leap_model_parser-0.1.250.dev4 → leap_model_parser-0.1.250.dev6}/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]:
@@ -74,6 +74,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.dev4 → leap_model_parser-0.1.250.dev6}/pyproject.toml

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