neurarch-trace 0.1.0__tar.gz

neurarch_trace-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 neurarch-ai
+
+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.

neurarch_trace-0.1.0/PKG-INFO

@@ -0,0 +1,124 @@
+Metadata-Version: 2.4
+Name: neurarch-trace
+Version: 0.1.0
+Summary: Run one forward pass over a PyTorch model and write a .neurarch.json graph with real per-layer shapes, for neurarch-mcp.
+Author: neurarch-ai
+License: MIT
+Project-URL: Homepage, https://github.com/neurarch-ai/neurarch-mcp
+Project-URL: Repository, https://github.com/neurarch-ai/neurarch-mcp
+Keywords: pytorch,mcp,neurarch,model-graph,tracing
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Provides-Extra: hf
+Requires-Dist: transformers; extra == "hf"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Dynamic: license-file
+
+# neurarch-trace
+
+Run one forward pass over a PyTorch model and write a `.neurarch.json` graph with
+the real input and output shape of every layer. Point
+[neurarch-mcp](https://github.com/neurarch-ai/neurarch-mcp) at that file and every
+tool works: parameter counts, FLOPs, shape contracts, `lint_model`, `check_design`.
+
+neurarch-mcp can already read a `.py` file by parsing it statically. Static parsing
+cannot see tensor shapes, and it cannot follow a model that is built at runtime:
+`AutoModel.from_pretrained(...)`, a timm factory, an architecture spread across a
+dozen files. `neurarch-trace` closes that gap by instantiating the model and
+watching the tensors go through it.
+
+## Install
+
+```
+pip install neurarch-trace          # torch >= 2.0
+pip install 'neurarch-trace[hf]'    # adds transformers for hf: targets
+```
+
+## Usage
+
+```
+neurarch-trace <target> --input 1,3,224,224 [--input 1,128:long ...] [-o out.neurarch.json]
+python -m neurarch_trace <target> ...     # same thing
+```
+
+A file and a class or factory in it:
+
+```
+neurarch-trace models/resnet.py:ResNet18 --input 1,3,224,224 -o resnet18.neurarch.json
+```
+
+An importable module and a factory function (called with no arguments):
+
+```
+neurarch-trace my_pkg.model:build_model --input 1,3,224,224
+```
+
+A Hugging Face checkpoint (defaults to `--input 1,16 --dtype long`, token ids):
+
+```
+neurarch-trace hf:prajjwal1/bert-tiny -o bert-tiny.neurarch.json
+```
+
+Then hand the graph to your agent:
+
+```
+npx -y neurarch-mcp ./resnet18.neurarch.json
+```
+
+`<target>` may name an `nn.Module` instance, an `nn.Module` subclass (instantiated
+with no arguments), or a callable that returns one. `--input` is repeated for
+multi-input forwards; a `:dtype` suffix (`1,128:long`) overrides `--dtype` for that
+input. Random tensors are used, `torch.randn` for float dtypes and
+`torch.randint(0, 1000, ...)` for integer ones.
+
+Other flags: `--name` sets the graph name (default: the attribute or repo name),
+`-o -` writes to stdout, `--depth N` stops descending at module depth N and
+records the modules there as single nodes, `--verbose` shows the traceback on a
+failure (otherwise a failure is one line on stderr and exit code 1).
+
+## The shape convention
+
+Shapes in the graph are written **without the batch dimension**: `[3, 224, 224]`
+for an image, `[128, 768]` for a token sequence. The first dim of every `--input`
+is the batch and is stripped from every recorded shape. Neurarch reads a leading
+dimension as the channel axis, so a shape that still carries its batch of 1 would
+be read as a one-channel tensor and every downstream number would be wrong.
+
+The trace runs in `eval()` mode on CPU with autograd enabled.
+
+## What ends up in the graph
+
+- One node per leaf module call, typed with the Neurarch vocabulary (`conv2d`,
+  `linear`, `layerNorm`, `multiHeadAttention`, ...) and carrying the same
+  parameter keys the static parser emits, so the MCP estimators read them. A
+  module called twice is two nodes. Modules the mapping table does not know
+  become `customModule` with `className` and `paramCount` in their params.
+- `scope` on every node is the dotted path of its parent module
+  (`encoder.layer.0.attention`), which is what the MCP block tools group by.
+- Edges follow actual data flow. A tensor a layer returns is matched by identity;
+  anything produced by functional code in between (`x + residual`, `torch.cat`,
+  reshapes, RoPE math) is traced back through autograd to the layers it came from.
+  A residual add becomes an `add` node with two inputs, a concatenation a
+  `concatenate` node.
+- Exactly one `input` node per `--input` and one `output` node.
+
+## Limits
+
+- Functional ops other than add / cat / stack do not get nodes of their own; a
+  `torch.flatten` or `F.softmax` between two modules is invisible, and the two
+  modules are wired directly.
+- Attribution runs on the autograd graph, so a forward wrapped in
+  `torch.no_grad()` or an integer-only path with no learnable parameter behind it
+  falls back to matching by dtype against the model inputs.
+- The batch is assumed to be dim 0. Sequence-first layouts
+  (`nn.MultiheadAttention` with `batch_first=False`) come out with the sequence
+  length stripped instead.
+- The mapping table lives in `neurarch_trace/mapping.py` and mirrors
+  `codeParser.ts` in the Neurarch app. If a torch module is missing there, it is
+  recorded as `customModule` rather than guessed.

neurarch_trace-0.1.0/README.md

@@ -0,0 +1,102 @@
+# neurarch-trace
+
+Run one forward pass over a PyTorch model and write a `.neurarch.json` graph with
+the real input and output shape of every layer. Point
+[neurarch-mcp](https://github.com/neurarch-ai/neurarch-mcp) at that file and every
+tool works: parameter counts, FLOPs, shape contracts, `lint_model`, `check_design`.
+
+neurarch-mcp can already read a `.py` file by parsing it statically. Static parsing
+cannot see tensor shapes, and it cannot follow a model that is built at runtime:
+`AutoModel.from_pretrained(...)`, a timm factory, an architecture spread across a
+dozen files. `neurarch-trace` closes that gap by instantiating the model and
+watching the tensors go through it.
+
+## Install
+
+```
+pip install neurarch-trace          # torch >= 2.0
+pip install 'neurarch-trace[hf]'    # adds transformers for hf: targets
+```
+
+## Usage
+
+```
+neurarch-trace <target> --input 1,3,224,224 [--input 1,128:long ...] [-o out.neurarch.json]
+python -m neurarch_trace <target> ...     # same thing
+```
+
+A file and a class or factory in it:
+
+```
+neurarch-trace models/resnet.py:ResNet18 --input 1,3,224,224 -o resnet18.neurarch.json
+```
+
+An importable module and a factory function (called with no arguments):
+
+```
+neurarch-trace my_pkg.model:build_model --input 1,3,224,224
+```
+
+A Hugging Face checkpoint (defaults to `--input 1,16 --dtype long`, token ids):
+
+```
+neurarch-trace hf:prajjwal1/bert-tiny -o bert-tiny.neurarch.json
+```
+
+Then hand the graph to your agent:
+
+```
+npx -y neurarch-mcp ./resnet18.neurarch.json
+```
+
+`<target>` may name an `nn.Module` instance, an `nn.Module` subclass (instantiated
+with no arguments), or a callable that returns one. `--input` is repeated for
+multi-input forwards; a `:dtype` suffix (`1,128:long`) overrides `--dtype` for that
+input. Random tensors are used, `torch.randn` for float dtypes and
+`torch.randint(0, 1000, ...)` for integer ones.
+
+Other flags: `--name` sets the graph name (default: the attribute or repo name),
+`-o -` writes to stdout, `--depth N` stops descending at module depth N and
+records the modules there as single nodes, `--verbose` shows the traceback on a
+failure (otherwise a failure is one line on stderr and exit code 1).
+
+## The shape convention
+
+Shapes in the graph are written **without the batch dimension**: `[3, 224, 224]`
+for an image, `[128, 768]` for a token sequence. The first dim of every `--input`
+is the batch and is stripped from every recorded shape. Neurarch reads a leading
+dimension as the channel axis, so a shape that still carries its batch of 1 would
+be read as a one-channel tensor and every downstream number would be wrong.
+
+The trace runs in `eval()` mode on CPU with autograd enabled.
+
+## What ends up in the graph
+
+- One node per leaf module call, typed with the Neurarch vocabulary (`conv2d`,
+  `linear`, `layerNorm`, `multiHeadAttention`, ...) and carrying the same
+  parameter keys the static parser emits, so the MCP estimators read them. A
+  module called twice is two nodes. Modules the mapping table does not know
+  become `customModule` with `className` and `paramCount` in their params.
+- `scope` on every node is the dotted path of its parent module
+  (`encoder.layer.0.attention`), which is what the MCP block tools group by.
+- Edges follow actual data flow. A tensor a layer returns is matched by identity;
+  anything produced by functional code in between (`x + residual`, `torch.cat`,
+  reshapes, RoPE math) is traced back through autograd to the layers it came from.
+  A residual add becomes an `add` node with two inputs, a concatenation a
+  `concatenate` node.
+- Exactly one `input` node per `--input` and one `output` node.
+
+## Limits
+
+- Functional ops other than add / cat / stack do not get nodes of their own; a
+  `torch.flatten` or `F.softmax` between two modules is invisible, and the two
+  modules are wired directly.
+- Attribution runs on the autograd graph, so a forward wrapped in
+  `torch.no_grad()` or an integer-only path with no learnable parameter behind it
+  falls back to matching by dtype against the model inputs.
+- The batch is assumed to be dim 0. Sequence-first layouts
+  (`nn.MultiheadAttention` with `batch_first=False`) come out with the sequence
+  length stripped instead.
+- The mapping table lives in `neurarch_trace/mapping.py` and mirrors
+  `codeParser.ts` in the Neurarch app. If a torch module is missing there, it is
+  recorded as `customModule` rather than guessed.

neurarch_trace-0.1.0/neurarch_trace/__init__.py

@@ -0,0 +1,19 @@
+"""neurarch-trace: one forward pass, one .neurarch.json with real shapes.
+
+    from neurarch_trace import trace_model
+    graph = trace_model(model, [torch.randn(1, 3, 224, 224)], name="resnet")
+"""
+from typing import Any, Dict, Optional, Sequence
+
+__version__ = "0.1.0"
+
+
+def trace_model(model, inputs: Sequence[Any], name: str = "model", depth: Optional[int] = None,
+                description: str = "") -> Dict[str, Any]:
+    """Trace `model` on `inputs` (batch first) and return the graph as a dict."""
+    from .tracer import trace
+    from .writer import build_graph
+    return build_graph(trace(model, list(inputs), depth=depth), name, description)
+
+
+__all__ = ["trace_model", "__version__"]

neurarch_trace-0.1.0/neurarch_trace/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+from .cli import main
+
+sys.exit(main())

neurarch_trace-0.1.0/neurarch_trace/cli.py

@@ -0,0 +1,155 @@
+"""`neurarch-trace <target> --input 1,3,224,224 [-o out.neurarch.json]`."""
+import argparse
+import importlib
+import importlib.util
+import os
+import re
+import sys
+from typing import List, Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+
+from . import __version__
+from .tracer import trace
+from .writer import build_graph, write_graph
+
+DTYPES = {
+    "float32": torch.float32, "float": torch.float32, "float16": torch.float16, "half": torch.float16,
+    "bfloat16": torch.bfloat16, "float64": torch.float64, "long": torch.long, "int64": torch.long,
+    "int": torch.int32, "int32": torch.int32, "bool": torch.bool,
+}
+
+
+class TraceError(Exception):
+    """A failure the user can act on; printed as one line without a traceback."""
+
+
+def load_target(target: str) -> Tuple[nn.Module, str]:
+    """Resolve `module:attr`, `file.py:attr` or `hf:<repo>` into a model and a default name."""
+    if target.startswith("hf:"):
+        repo = target[3:]
+        try:
+            from transformers import AutoModel
+        except ImportError:
+            raise TraceError("hf: targets need transformers: pip install 'neurarch-trace[hf]'")
+        return AutoModel.from_pretrained(repo), re.sub(r"[^A-Za-z0-9_.-]+", "-", repo.split("/")[-1])
+
+    if ":" not in target:
+        raise TraceError("target must be module.path:attr, path/to/file.py:attr or hf:<repo-id>, got %r" % target)
+    modpath, attr = target.rsplit(":", 1)
+    if modpath.endswith(".py") or os.sep in modpath:
+        path = os.path.abspath(modpath)
+        if not os.path.isfile(path):
+            raise TraceError("no such file: %s" % path)
+        sys.path.insert(0, os.path.dirname(path))
+        spec = importlib.util.spec_from_file_location(os.path.splitext(os.path.basename(path))[0], path)
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+    else:
+        # A bare module path is resolved from the working directory first, the
+        # way `python -m` would, so `models.resnet:ResNet18` works from a repo root.
+        sys.path.insert(0, os.getcwd())
+        mod = importlib.import_module(modpath)
+
+    obj = mod
+    for part in attr.split("."):
+        if not hasattr(obj, part):
+            raise TraceError("%s has no attribute %r" % (modpath, attr))
+        obj = getattr(obj, part)
+
+    if isinstance(obj, nn.Module):
+        model = obj
+    elif callable(obj):
+        try:
+            model = obj()
+        except Exception as e:
+            raise TraceError("%s() failed: %s: %s" % (attr, type(e).__name__, e))
+        if not isinstance(model, nn.Module):
+            raise TraceError("%s() returned %s, not an nn.Module" % (attr, type(model).__name__))
+    else:
+        raise TraceError("%s is %s; expected an nn.Module, an nn.Module subclass or a factory" % (attr, type(obj).__name__))
+    return model, attr.split(".")[-1]
+
+
+def build_inputs(specs: Sequence[str], default_dtype: str) -> List[torch.Tensor]:
+    """`1,3,224,224` or `1,128:long`; the first dim is the batch."""
+    torch.manual_seed(0)
+    out = []
+    for spec in specs:
+        dims_text, _, dtype_name = spec.partition(":")
+        dtype_name = dtype_name or default_dtype
+        if dtype_name not in DTYPES:
+            raise TraceError("unknown dtype %r (one of %s)" % (dtype_name, ", ".join(sorted(DTYPES))))
+        try:
+            dims = [int(d) for d in dims_text.split(",") if d.strip()]
+        except ValueError:
+            raise TraceError("bad --input %r: expected comma-separated integers like 1,3,224,224" % spec)
+        if not dims:
+            raise TraceError("bad --input %r: no dims" % spec)
+        dtype = DTYPES[dtype_name]
+        if dtype == torch.bool:
+            out.append(torch.rand(dims) > 0.5)
+        elif dtype.is_floating_point:
+            out.append(torch.randn(dims, dtype=dtype))
+        else:
+            out.append(torch.randint(0, 1000, dims, dtype=dtype))
+    return out
+
+
+def run(argv: Optional[Sequence[str]] = None) -> int:
+    p = argparse.ArgumentParser(
+        prog="neurarch-trace",
+        description="Run one forward pass over a PyTorch model and write a .neurarch.json graph with real shapes.",
+    )
+    p.add_argument("target", help="module.path:attr, path/to/file.py:attr, or hf:<repo-id>")
+    p.add_argument("--input", action="append", default=[], metavar="DIMS",
+                   help="input dims, batch first, e.g. 1,3,224,224 or 1,128:long; repeat for multi-input forwards")
+    p.add_argument("--dtype", default=None, help="dtype for inputs without a :dtype suffix (default float32; long for hf:)")
+    p.add_argument("--name", default=None, help="graph name (default: the attr or repo name)")
+    p.add_argument("-o", "--output", default=None, help="output path (default ./<name>.neurarch.json; - for stdout)")
+    p.add_argument("--depth", type=int, default=None, help="stop descending at this module depth (default: leaf modules)")
+    p.add_argument("--verbose", action="store_true", help="show the full traceback on failure")
+    p.add_argument("--version", action="version", version="neurarch-trace " + __version__)
+    args = p.parse_args(argv)
+
+    is_hf = args.target.startswith("hf:")
+    specs = args.input or (["1,16"] if is_hf else [])
+    if not specs:
+        raise TraceError("--input is required (e.g. --input 1,3,224,224); it defaults only for hf: targets")
+    dtype = args.dtype or ("long" if is_hf else "float32")
+
+    model, default_name = load_target(args.target)
+    inputs = build_inputs(specs, dtype)
+    try:
+        nodes = trace(model, inputs, depth=args.depth)
+    except TraceError:
+        raise
+    except Exception as e:
+        raise TraceError("forward pass failed: %s: %s" % (type(e).__name__, e))
+
+    name = args.name or default_name
+    description = "Traced by neurarch-trace %s from %s with input %s" % (
+        __version__, args.target, " ".join(specs))
+    graph = build_graph(nodes, name, description)
+    out = args.output or ("./%s.neurarch.json" % name)
+    write_graph(graph, out)
+    if out != "-":
+        layers = sum(1 for c in graph["components"] if c["type"] not in ("input", "output"))
+        print("wrote %s (%d layers, %d connections)" % (out, layers, len(graph["connections"])))
+    return 0
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    verbose = "--verbose" in (argv if argv is not None else sys.argv[1:])
+    try:
+        return run(argv)
+    except Exception as e:
+        if verbose:
+            raise
+        print("neurarch-trace: %s" % e, file=sys.stderr)
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

neurarch_trace-0.1.0/neurarch_trace/mapping.py

@@ -0,0 +1,203 @@
+"""torch module class -> Neurarch component type and param keys.
+
+Source of truth: `src/utils/codeParser.ts` in the Neurarch app (the static
+parser that neurarch-mcp vendors). The MCP param and FLOPs estimators key off
+the exact param names that parser emits, so every entry here mirrors what
+`parseLayerDefinition` would produce for the same `nn.X(...)` call. Adding a
+type that codeParser does not emit is fine as long as it exists in
+`src/lib/types.ts`; renaming a param key is not.
+"""
+from typing import Any, Callable, Dict, Optional, Tuple
+
+import torch.nn as nn
+
+Params = Dict[str, Any]
+
+
+def _first(v: Any) -> Any:
+    # torch stores kernel_size / stride / padding as tuples; the graph stores the
+    # scalar codeParser reads off the source when all dims agree.
+    if isinstance(v, (tuple, list)):
+        return v[0] if all(x == v[0] for x in v) else list(v)
+    return v
+
+
+def _conv(m: nn.Module) -> Params:
+    p: Params = {
+        "inChannels": m.in_channels,
+        "outChannels": m.out_channels,
+        "kernelSize": _first(m.kernel_size),
+        "stride": _first(m.stride),
+        "padding": _first(m.padding),
+    }
+    if _first(m.dilation) != 1:
+        p["dilation"] = _first(m.dilation)
+    if m.groups != 1:
+        p["groups"] = m.groups
+    if m.bias is None:
+        p["bias"] = False
+    return p
+
+
+def _conv2d(m: nn.Module) -> Tuple[str, Params]:
+    # codeParser's ViT heuristic: a wide conv whose stride equals its kernel on
+    # a few input channels is a patch projection, not a feature extractor.
+    k, s = _first(m.kernel_size), _first(m.stride)
+    if isinstance(k, int) and k > 4 and k == s and m.in_channels <= 4:
+        return "patchEmbed", {"patchSize": k, "embedDim": m.out_channels, "inChans": m.in_channels}
+    return "conv2d", _conv(m)
+
+
+def _conv_t2d(m: nn.Module) -> Params:
+    p = _conv(m)
+    if _first(m.output_padding) != 0:
+        p["outputPadding"] = _first(m.output_padding)
+    return p
+
+
+def _pool(m: nn.Module) -> Params:
+    return {"kernelSize": _first(m.kernel_size), "stride": _first(m.stride or m.kernel_size)}
+
+
+def _adaptive_avg(m: nn.Module) -> Tuple[str, Params]:
+    out = m.output_size if isinstance(m.output_size, (tuple, list)) else (m.output_size,)
+    if all(o in (1, None) for o in out):
+        return "globalAvgPool2d", {}
+    return "adaptiveAvgPool2d", {"outputSize": _first(m.output_size)}
+
+
+def _recurrent(m: nn.Module) -> Params:
+    p: Params = {"inputSize": m.input_size, "hiddenSize": m.hidden_size, "numLayers": m.num_layers}
+    if m.bidirectional:
+        p["bidirectional"] = True
+    return p
+
+
+def _norm_shape(m: nn.Module) -> Params:
+    return {"normalizedShape": _first(tuple(m.normalized_shape))}
+
+
+# class name -> (component type, param extractor). An extractor may instead
+# return (type, params) itself when the type depends on the module's config.
+TORCH_CLASS_MAP: Dict[str, Tuple[Optional[str], Callable[[nn.Module], Any]]] = {
+    "Linear": ("linear", lambda m: dict(
+        {"inFeatures": m.in_features, "outFeatures": m.out_features}, **({} if m.bias is not None else {"bias": False}))),
+    "Conv2d": (None, _conv2d),
+    "ConvTranspose2d": ("transposeConv2d", _conv_t2d),
+    "Conv1d": ("conv1d", _conv),
+    "Conv3d": ("conv3d", _conv),
+    "MaxPool2d": ("maxpool2d", _pool),
+    "AvgPool2d": ("avgpool2d", _pool),
+    "MaxPool1d": ("maxpool1d", _pool),
+    "AvgPool1d": ("avgpool1d", _pool),
+    "AdaptiveAvgPool2d": (None, _adaptive_avg),
+    "AdaptiveMaxPool2d": ("adaptiveMaxPool2d", lambda m: {"outputSize": _first(m.output_size)}),
+    "Upsample": ("upsample", lambda m: {"scaleFactor": _first(m.scale_factor)}),
+    "Dropout": ("dropout", lambda m: {"p": m.p}),
+    "Dropout2d": ("dropout", lambda m: {"p": m.p}),
+    "BatchNorm1d": ("batchNorm", lambda m: {"numFeatures": m.num_features}),
+    "BatchNorm2d": ("batchNorm", lambda m: {"numFeatures": m.num_features}),
+    "BatchNorm3d": ("batchNorm", lambda m: {"numFeatures": m.num_features}),
+    "LayerNorm": ("layerNorm", _norm_shape),
+    "RMSNorm": ("rmsNorm", _norm_shape),
+    "GroupNorm": ("groupNorm", lambda m: {"numGroups": m.num_groups, "numChannels": m.num_channels}),
+    "InstanceNorm2d": ("instanceNorm", lambda m: {"numFeatures": m.num_features}),
+    "Embedding": ("embedding", lambda m: {"vocabSize": m.num_embeddings, "embeddingDim": m.embedding_dim}),
+    "EmbeddingBag": ("embeddingBag", lambda m: {"vocabSize": m.num_embeddings, "embeddingDim": m.embedding_dim}),
+    "MultiheadAttention": ("multiHeadAttention", lambda m: {"hiddenDim": m.embed_dim, "numHeads": m.num_heads}),
+    "TransformerEncoderLayer": ("transformerBlock", lambda m: {
+        "embedDim": m.self_attn.embed_dim, "numHeads": m.self_attn.num_heads, "ffDim": m.linear1.out_features}),
+    "LSTM": ("lstm", _recurrent),
+    "GRU": ("gru", _recurrent),
+    "RNN": ("rnn", _recurrent),
+    "Flatten": ("flatten", lambda m: {}),
+    "ReLU": ("relu", lambda m: {}),
+    "ReLU6": ("relu6", lambda m: {}),
+    "LeakyReLU": ("leakyRelu", lambda m: {"negativeSlope": m.negative_slope}),
+    "ELU": ("elu", lambda m: {}),
+    "PReLU": ("prelu", lambda m: {"numParameters": m.num_parameters}),
+    "SELU": ("selu", lambda m: {}),
+    "GELU": ("gelu", lambda m: {}),
+    "SiLU": ("swish", lambda m: {}),
+    "Mish": ("mish", lambda m: {}),
+    "Hardswish": ("hardSwish", lambda m: {}),
+    "Hardsigmoid": ("hardSigmoid", lambda m: {}),
+    "Sigmoid": ("sigmoid", lambda m: {}),
+    "Tanh": ("tanh", lambda m: {}),
+    "Softmax": ("softmax", lambda m: {}),
+    "LogSoftmax": ("logSoftmax", lambda m: {}),
+    "Softplus": ("softplus", lambda m: {}),
+    "GLU": ("glu", lambda m: {}),
+}
+
+# Modules that pass their input through untouched. They get no node, so the
+# tensor keeps pointing at its real producer and no fake hop appears.
+PASSTHROUGH_CLASSES = {"Identity"}
+
+# Copied from codeParser.ts CUSTOM_CLASS_MAP (lowercased class name -> type).
+# Only consulted when --depth stops descent at a container the torch table
+# does not know, so a `LlamaAttention` cut off at depth 3 is still an
+# attention node rather than an anonymous custom box.
+CUSTOM_CLASS_MAP: Dict[str, str] = {
+    "positionalencoding": "positionalEncoding", "posencoding": "positionalEncoding",
+    "posemb": "positionalEncoding", "positionalembedding": "positionalEncoding",
+    "rotaryembedding": "rope", "ropeembedding": "rope",
+    "multiheadattention": "multiHeadAttention", "multiheadattn": "multiHeadAttention",
+    "causalselfattention": "causalAttention", "selfattn": "selfAttention",
+    "selfattention": "selfAttention", "attention": "attention", "crossattention": "crossModalAttention",
+    "llamaattention": "groupedQueryAttention", "mistralattention": "groupedQueryAttention",
+    "mixtralattention": "groupedQueryAttention", "qwenattention": "groupedQueryAttention",
+    "gemmaattention": "groupedQueryAttention", "phi3attention": "groupedQueryAttention",
+    "falcon7battention": "groupedQueryAttention", "groupedqueryattention": "groupedQueryAttention",
+    "gqaattention": "groupedQueryAttention",
+    "feedforward": "feedForward", "feedforwardnetwork": "feedForward", "ffn": "feedForward", "mlp": "feedForward",
+    "llamamlp": "swiglu", "mistralmpl": "swiglu", "mixtralmlp": "swiglu", "qwenmlp": "swiglu",
+    "gemmamlp": "swiglu", "phi3mlp": "swiglu", "swiglu": "swiglu", "gatedmlp": "swiglu",
+    "mixtralsparsemoeblock": "moeLayer", "moelayer": "moeLayer", "expertlayer": "moeLayer", "sparsemlp": "moeLayer",
+    "block": "transformerBlock", "gptblock": "transformerBlock", "bertlayer": "transformerBlock",
+    "bertblock": "transformerBlock", "encoderlayer": "transformerBlock", "decoderlayer": "transformerBlock",
+    "transformerblock": "transformerBlock", "transformerlayer": "transformerBlock",
+    "transformer": "transformerBlock", "visionblock": "transformerBlock",
+    "llamadecoderlayer": "transformerBlock", "llamadecoderblock": "transformerBlock",
+    "mistraldecoderlayer": "transformerBlock", "mixtraldecoderlayer": "transformerBlock",
+    "qwendecoderlayer": "transformerBlock", "gemmadecoderlayer": "transformerBlock",
+    "phi3decoderlayer": "transformerBlock",
+    "visiontransformer": "transformerBlock", "vitblock": "transformerBlock", "vitlayer": "transformerBlock",
+    "patchembed": "patchEmbed", "patchembedding": "patchEmbed", "patchprojection": "patchEmbed",
+    "seblock": "seBlock", "squeezeexcitation": "seBlock", "channelattention": "seBlock",
+    "resblock": "residual", "residualblock": "residual", "resnetblock": "residual", "bottleneck": "residual",
+}
+
+# Attribute names real-world blocks use for the two numbers the estimators
+# need most. Read off a container only when it is cut off by --depth.
+_DIM_ATTRS = ("embed_dim", "d_model", "hidden_size", "dim", "n_embd")
+_HEAD_ATTRS = ("num_heads", "n_head", "nhead", "num_attention_heads")
+
+
+def _sniff(m: nn.Module, names: Tuple[str, ...]) -> Optional[int]:
+    for n in names:
+        v = getattr(m, n, None)
+        if isinstance(v, int) and not isinstance(v, bool):
+            return v
+    return None
+
+
+def map_module(m: nn.Module) -> Tuple[str, Params]:
+    """Return (component type, params) for one module treated as a graph node."""
+    cls = type(m).__name__
+    entry = TORCH_CLASS_MAP.get(cls)
+    if entry is not None:
+        typ, extract = entry
+        result = extract(m)
+        return (typ, result) if typ is not None else result
+
+    n_params = sum(p.numel() for p in m.parameters())
+    params: Params = {"className": cls, "paramCount": n_params}
+    typ = CUSTOM_CLASS_MAP.get(cls.lower(), "customModule")
+    if typ != "customModule":
+        d, h = _sniff(m, _DIM_ATTRS), _sniff(m, _HEAD_ATTRS)
+        if d is not None:
+            params["embedDim"] = d
+        if h is not None:
+            params["numHeads"] = h
+    return typ, params