nnschemaparser 0.1.0__tar.gz

nnschemaparser-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Cerisara Nathan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

nnschemaparser-0.1.0/PKG-INFO

@@ -0,0 +1,200 @@
+Metadata-Version: 2.4
+Name: nnschemaparser
+Version: 0.1.0
+Summary: Parse torch nn.Module models into an inspectable ModelBlockSchema graph and execute it with a faithful schema interpreter
+Author: Cerisara Nathan
+License: MIT License
+        
+        Copyright (c) 2025 Cerisara Nathan
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://github.com/nath54/NNViewer
+Project-URL: Documentation, https://github.com/nath54/NNViewer/blob/main/docs/DOCUMENTATION.md
+Project-URL: Changelog, https://github.com/nath54/NNViewer/blob/main/CHANGELOG.md
+Project-URL: Issues, https://github.com/nath54/NNViewer/issues
+Keywords: pytorch,neural-network,graph,schema,interpreter,model-inspection
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Provides-Extra: research
+Requires-Dist: transformers>=4.40; extra == "research"
+Requires-Dist: espnet; extra == "research"
+Dynamic: license-file
+
+# NNSchemaParser
+
+**NNSchemaParser** converts a PyTorch `nn.Module` into a serializable, inspectable
+graph representation — the **`ModelBlockSchema`** — that keeps the *original variable
+names of the forward source* and *every sub-module* (including MoE experts that a
+single trace would never visit). The schema can then be executed by a built-in
+**schema interpreter** that re-dispatches every recorded call to the real torch
+callables, so the graph runs exactly like the model — and can optionally capture
+the value of every intermediate tensor.
+
+```
+   nn.Module  ──parse──►  ModelBlockSchema  ──interpret──►  outputs / captured values
+  (your code)              (names + graph)                   (identical to model(x))
+```
+
+## Why not `torch.fx`?
+
+| | `torch.fx.symbolic_trace` | **NNSchemaParser** |
+|---|---|---|
+| Variable names | lost (`_0`, `_1`, …) | **kept** (`hidden_states`, `query`, …) |
+| Untaken branches (MoE experts) | missing | **all present** |
+| Control flow | `graph_module` nodes only | `if` / `for` / `try` as first-class instructions |
+| `*args` / `**kwargs` | limited | supported in signatures and call sites |
+| Comprehensions | unsupported | compiled to helper functions |
+| Re-executable | via the FX machinery | via a small self-contained interpreter |
+
+## Quick start
+
+```python
+import torch
+from torch import nn
+
+from nnschemaparser import NNParser, SchemaInterpreter, TraceRecorder, dump_schema
+
+
+class MLP(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(784, 256)
+        self.fc2 = nn.Linear(256, 10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        hidden = self.fc1(x)
+        activated = hidden.relu()
+        return self.fc2(activated)
+
+
+model = MLP().eval()
+x = torch.randn(2, 784)
+
+# 1. Parse the model into a schema (resolves shapes and dtypes)
+schema = NNParser.parse_model(
+    model=model,
+    inputs_tensors_shapes={"x": (2, 784)},
+    inputs_tensors_dtypes={"x": "torch.float32"},
+    inputs_tensors_values={"x": x},
+)
+
+# 2. Dump a readable view of the whole block tree
+dump_schema(schema=schema, path="schema.txt")
+
+# 3. Execute it with a tracer and compare with the real forward
+tracer = TraceRecorder()
+outputs = SchemaInterpreter(capture_values=True, listeners=[tracer]).run(
+    block=schema,
+    scope={"x": x.clone()},
+)
+tracer.save("trace.txt")
+reference = model(x)
+assert torch.equal(outputs[0], reference)
+```
+
+The schema dump of `MLP` reads exactly like the source — original variable
+names, every sub-module, and the resolved shape of every step:
+
+```
+=== BLOCK MLP ===
+inputs:
+  - x = torch.float32[2, 784]
+outputs:
+  - torch.float32[2, 10]
+weights:
+  (none)
+sub-blocks:
+  - fc1(Linear)
+  - fc2(Linear)
+body:
+  hidden = CALL fc1(x)  -> torch.float32[2, 256]
+  $t1 = getattr(hidden, 'relu')  -> builtin_function_or_method
+  activated = hidden.relu()  -> torch.float32[2, 256]
+  $t2 = CALL fc2(activated)  -> torch.float32[2, 10]
+  RETURN $t2
+```
+
+The tracer records the same execution chronologically, with the value of
+every intermediate tensor (one line per event, indented by the nesting
+depth) — this is what a WebUI consumes to display every intermediate value:
+
+```
+[00001] ENTER BLOCK MLP
+[00002]   in x = torch.float32[2, 784] abs_sum=1241.4496 min=-3.0303 max=3.3046
+[00003]   ENTER BLOCK Linear
+[00004]     in input = torch.float32[2, 784] abs_sum=1241.4496 min=-3.0303 max=3.3046
+[00005]     $t1 = F.linear(input, weight, bias)  [torch.float32[2, 256] abs_sum=215.3293 ...]
+[00006]   EXIT BLOCK Linear -> (torch.float32[2, 256] abs_sum=215.3293 ...)
+[00007]   hidden = CALL fc1 [Linear]  [torch.float32[2, 256] abs_sum=215.3293 ...]
+```
+
+## Installation
+
+```bash
+pip install .            # or: pip install -e .[research] for the harness extras
+```
+
+Requires Python ≥ 3.12 and PyTorch ≥ 2.0. The model classes must be importable
+from a **real file** (the parser reads the forward source with
+`inspect.getsource`, which fails for classes defined in a REPL or stdin).
+
+## Validation
+
+A robustness sweep of **58 models** — 48 Hugging Face hub checkpoints (text,
+vision, audio, multimodal) plus 10 locally built architectures (MLP, RNN, CNN,
+autoencoder, GAN, GNN) — runs each model through *parse → interpret → compare
+against the real forward*. **All 58 pass with `maxdiff = 0.0`** against the real
+forward. See [docs/DOCUMENTATION.md](docs/DOCUMENTATION.md) for the full model
+matrix and the sweep CLI.
+
+## Documentation
+
+- **[docs/API.md](docs/API.md)** — the public API: parsing, executing,
+  tracing, dumping and walking schemas.
+- **[docs/DOCUMENTATION.md](docs/DOCUMENTATION.md)** — full documentation:
+  how the parser works, the schema and instruction reference, the
+  interpreter, usage recipes, supported syntax, supported/unsupported
+  models, and limitations.
+- **[CHANGELOG.md](CHANGELOG.md)** — the release history.
+- `tests/hf_sweep.py` — the validation harness CLI (58 HuggingFace models).
+- `tests/test_new_syntax.py`, `tests/test_parser.py`,
+  `tests/test_methods_and_scalars.py` — the fixture test suite (pytest,
+  no network needed), also runnable with `python -m tests.test_new_syntax`.
+- `tests/test_speech_models.py` — the espnet/torchaudio regression tests.
+
+## License
+
+MIT — see [LICENSE](LICENSE).

nnschemaparser-0.1.0/README.md

@@ -0,0 +1,145 @@
+# NNSchemaParser
+
+**NNSchemaParser** converts a PyTorch `nn.Module` into a serializable, inspectable
+graph representation — the **`ModelBlockSchema`** — that keeps the *original variable
+names of the forward source* and *every sub-module* (including MoE experts that a
+single trace would never visit). The schema can then be executed by a built-in
+**schema interpreter** that re-dispatches every recorded call to the real torch
+callables, so the graph runs exactly like the model — and can optionally capture
+the value of every intermediate tensor.
+
+```
+   nn.Module  ──parse──►  ModelBlockSchema  ──interpret──►  outputs / captured values
+  (your code)              (names + graph)                   (identical to model(x))
+```
+
+## Why not `torch.fx`?
+
+| | `torch.fx.symbolic_trace` | **NNSchemaParser** |
+|---|---|---|
+| Variable names | lost (`_0`, `_1`, …) | **kept** (`hidden_states`, `query`, …) |
+| Untaken branches (MoE experts) | missing | **all present** |
+| Control flow | `graph_module` nodes only | `if` / `for` / `try` as first-class instructions |
+| `*args` / `**kwargs` | limited | supported in signatures and call sites |
+| Comprehensions | unsupported | compiled to helper functions |
+| Re-executable | via the FX machinery | via a small self-contained interpreter |
+
+## Quick start
+
+```python
+import torch
+from torch import nn
+
+from nnschemaparser import NNParser, SchemaInterpreter, TraceRecorder, dump_schema
+
+
+class MLP(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(784, 256)
+        self.fc2 = nn.Linear(256, 10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        hidden = self.fc1(x)
+        activated = hidden.relu()
+        return self.fc2(activated)
+
+
+model = MLP().eval()
+x = torch.randn(2, 784)
+
+# 1. Parse the model into a schema (resolves shapes and dtypes)
+schema = NNParser.parse_model(
+    model=model,
+    inputs_tensors_shapes={"x": (2, 784)},
+    inputs_tensors_dtypes={"x": "torch.float32"},
+    inputs_tensors_values={"x": x},
+)
+
+# 2. Dump a readable view of the whole block tree
+dump_schema(schema=schema, path="schema.txt")
+
+# 3. Execute it with a tracer and compare with the real forward
+tracer = TraceRecorder()
+outputs = SchemaInterpreter(capture_values=True, listeners=[tracer]).run(
+    block=schema,
+    scope={"x": x.clone()},
+)
+tracer.save("trace.txt")
+reference = model(x)
+assert torch.equal(outputs[0], reference)
+```
+
+The schema dump of `MLP` reads exactly like the source — original variable
+names, every sub-module, and the resolved shape of every step:
+
+```
+=== BLOCK MLP ===
+inputs:
+  - x = torch.float32[2, 784]
+outputs:
+  - torch.float32[2, 10]
+weights:
+  (none)
+sub-blocks:
+  - fc1(Linear)
+  - fc2(Linear)
+body:
+  hidden = CALL fc1(x)  -> torch.float32[2, 256]
+  $t1 = getattr(hidden, 'relu')  -> builtin_function_or_method
+  activated = hidden.relu()  -> torch.float32[2, 256]
+  $t2 = CALL fc2(activated)  -> torch.float32[2, 10]
+  RETURN $t2
+```
+
+The tracer records the same execution chronologically, with the value of
+every intermediate tensor (one line per event, indented by the nesting
+depth) — this is what a WebUI consumes to display every intermediate value:
+
+```
+[00001] ENTER BLOCK MLP
+[00002]   in x = torch.float32[2, 784] abs_sum=1241.4496 min=-3.0303 max=3.3046
+[00003]   ENTER BLOCK Linear
+[00004]     in input = torch.float32[2, 784] abs_sum=1241.4496 min=-3.0303 max=3.3046
+[00005]     $t1 = F.linear(input, weight, bias)  [torch.float32[2, 256] abs_sum=215.3293 ...]
+[00006]   EXIT BLOCK Linear -> (torch.float32[2, 256] abs_sum=215.3293 ...)
+[00007]   hidden = CALL fc1 [Linear]  [torch.float32[2, 256] abs_sum=215.3293 ...]
+```
+
+## Installation
+
+```bash
+pip install .            # or: pip install -e .[research] for the harness extras
+```
+
+Requires Python ≥ 3.12 and PyTorch ≥ 2.0. The model classes must be importable
+from a **real file** (the parser reads the forward source with
+`inspect.getsource`, which fails for classes defined in a REPL or stdin).
+
+## Validation
+
+A robustness sweep of **58 models** — 48 Hugging Face hub checkpoints (text,
+vision, audio, multimodal) plus 10 locally built architectures (MLP, RNN, CNN,
+autoencoder, GAN, GNN) — runs each model through *parse → interpret → compare
+against the real forward*. **All 58 pass with `maxdiff = 0.0`** against the real
+forward. See [docs/DOCUMENTATION.md](docs/DOCUMENTATION.md) for the full model
+matrix and the sweep CLI.
+
+## Documentation
+
+- **[docs/API.md](docs/API.md)** — the public API: parsing, executing,
+  tracing, dumping and walking schemas.
+- **[docs/DOCUMENTATION.md](docs/DOCUMENTATION.md)** — full documentation:
+  how the parser works, the schema and instruction reference, the
+  interpreter, usage recipes, supported syntax, supported/unsupported
+  models, and limitations.
+- **[CHANGELOG.md](CHANGELOG.md)** — the release history.
+- `tests/hf_sweep.py` — the validation harness CLI (58 HuggingFace models).
+- `tests/test_new_syntax.py`, `tests/test_parser.py`,
+  `tests/test_methods_and_scalars.py` — the fixture test suite (pytest,
+  no network needed), also runnable with `python -m tests.test_new_syntax`.
+- `tests/test_speech_models.py` — the espnet/torchaudio regression tests.
+
+## License
+
+MIT — see [LICENSE](LICENSE).

nnschemaparser-0.1.0/nnschemaparser/__init__.py

@@ -0,0 +1,27 @@
+"""
+NNSchemaParser.
+
+Converts a torch nn.Module into an inspectable ModelBlockSchema graph that
+keeps the original variable names of the forward source and every sub-module,
+and executes it with a schema interpreter re-dispatching to the real torch
+callables.
+"""
+
+# Import Modules
+from nnschemaparser.core import ModelBlockSchema, TensorRepr, TensorReprWithValue
+from nnschemaparser.dump import dump_schema
+from nnschemaparser.interpreter import SchemaInterpreter, TraceListener, TraceRecorder
+from nnschemaparser.parser import NNParser
+
+__all__ = [
+    "ModelBlockSchema",
+    "TraceListener",
+    "TraceRecorder",
+    "dump_schema",
+    "NNParser",
+    "SchemaInterpreter",
+    "TensorRepr",
+    "TensorReprWithValue"
+]
+
+__version__: str = "0.1.0"

nnschemaparser-0.1.0/nnschemaparser/core/__init__.py

@@ -0,0 +1,16 @@
+"""
+Core data model of NNSchemaParser.
+
+The schema types describing a parsed model: the block schema tree, the tensor
+representations and the instruction types of the block bodies.
+"""
+
+# Import Modules
+from nnschemaparser.core.model_schema import ModelBlockSchema
+from nnschemaparser.core.tensor_schema import TensorRepr, TensorReprWithValue
+
+__all__ = [
+    "ModelBlockSchema",
+    "TensorRepr",
+    "TensorReprWithValue"
+]

nnschemaparser-0.1.0/nnschemaparser/core/container_keys.py

@@ -0,0 +1,69 @@
+"""
+Container entry key utilities shared by the lowering and the interpreter.
+
+The entries of a sub-block container are keyed by "container.index" in the
+sub-block dictionaries, and must be iterated in their declaration order.
+"""
+
+# Import Modules
+
+
+def container_entry_sort_key(entry_key: str) -> int:
+    """
+    Get the numeric sort key of a container entry key.
+
+    Args:
+        entry_key (str): The entry key (e.g. "encoders.2").
+
+    Returns:
+        int: The numeric suffix of the key, or -1 when not numeric.
+    """
+
+    # Split the key into its container prefix and its suffix
+    parts: list[str] = entry_key.rsplit(".", 1)
+
+    # A numeric suffix is sorted by its value
+    if len(parts) == 2 and parts[1].isdigit():
+        return int(parts[1])
+
+    # A non numeric suffix is sorted last
+    return -1
+
+
+def container_entry_keys(
+    sub_block_names: set[str] | dict[str, str],
+    container_name: str
+) -> list[str]:
+    """
+    Get the ordered entry keys of a container in the parsed sub-blocks.
+
+    Args:
+        sub_block_names (set[str] | dict[str, str]): The names of the parsed
+            sub-blocks, either a set of names or the sub-block dictionary.
+        container_name (str): The name of the container.
+
+    Returns:
+        list[str]: The entry keys sorted by their numeric suffix.
+    """
+
+    # The special self path holds the direct numeric entries of the module
+    if container_name == "$self":
+        return sorted(
+            [
+                sub_block_name for sub_block_name in sub_block_names
+                if sub_block_name.isdigit()
+            ],
+            key=container_entry_sort_key
+        )
+
+    # Collect the sub-block keys of the container entries
+    keys: list[str] = [
+        sub_block_name for sub_block_name in sub_block_names
+        if sub_block_name.startswith(f"{container_name}.")
+    ]
+
+    # Sort the entry keys by their numeric suffix
+    keys.sort(key=container_entry_sort_key)
+
+    # Return the ordered entry keys of the container
+    return keys