metrik-torch 0.1.0__tar.gz

metrik_torch-0.1.0/.gitignore

@@ -0,0 +1,89 @@
+# Metrik — .gitignore
+#
+# NOTE: the previous version of this file contained `docs/*`, `.gitignore`, and `.claude`,
+# which excluded the entire documentation tree — i.e. every Phase 0/1 deliverable — from
+# version control. Replaced deliberately. docs/ is the project's primary output right now.
+
+# --- Build output that must never be committed ---
+# ADR 016: diagram sources live in docs/diagrams/*.mmd; rendered images are build output.
+docs/static/diagrams/
+# ADR 015: generated reference docs are committed on purpose (drift detection), so they
+# are NOT ignored here — see docs/reference/.
+
+# --- Python ---
+__pycache__/
+*.py[cod]
+*.so
+.Python
+build/
+dist/
+*.egg-info/
+.eggs/
+.venv/
+venv/
+env/
+
+# Tooling caches
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.hypothesis/
+.coverage
+.coverage.*
+htmlcov/
+coverage.xml
+.benchmarks/
+.tox/
+.nox/
+
+# uv (ADR 013) — uv.lock IS committed; the cache is not.
+.uv/
+
+# --- Node / web (ADR 009) ---
+node_modules/
+.next/
+out/
+.turbo/
+*.tsbuildinfo
+.pnpm-store/
+
+# --- Tauri (ADR 010) ---
+apps/desktop/src-tauri/target/
+apps/desktop/src-tauri/gen/
+
+# --- Docs site (ADR 015) ---
+docs/.docusaurus/
+docs/build/
+
+# --- Metrik runtime data ---
+# The artifact store and blob store (ADR 006). Content-addressed, machine-generated,
+# gigabytes. Never in git.
+.metrik/
+*.metrik-store/
+# Local config may hold API keys (ADR 018).
+metrik.local.toml
+.env
+.env.*
+!.env.example
+
+# Model weights and checkpoints — always too large, often license-encumbered (plan.md §5.7).
+*.safetensors
+*.gguf
+*.onnx
+*.pt
+*.pth
+*.bin
+# ...except the tiny-random fixtures, which are committed on purpose (ADR 017).
+!tests/fixtures/**
+
+# --- Editors and OS ---
+.vscode/
+!.vscode/extensions.json
+.idea/
+*.swp
+.DS_Store
+Thumbs.db
+desktop.ini
+
+# --- Local agent/tooling settings ---
+.claude/settings.local.json

metrik_torch-0.1.0/PKG-INFO

@@ -0,0 +1,49 @@
+Metadata-Version: 2.4
+Name: metrik-torch
+Version: 0.1.0
+Summary: Metrik torch loader — observed ModelGraph extraction from a real nn.Module tree.
+Project-URL: Homepage, https://github.com/Asmodeus14/Metrik
+Project-URL: Repository, https://github.com/Asmodeus14/Metrik
+Project-URL: Issues, https://github.com/Asmodeus14/Metrik/issues
+Author: The Metrik Authors
+License-Expression: Apache-2.0
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: metrik-sdk==0.1.0
+Provides-Extra: torch
+Requires-Dist: torch>=2.2; extra == 'torch'
+Description-Content-Type: text/markdown
+
+# metrik-torch
+
+Observed `ModelGraph` extraction from a real `nn.Module` tree.
+
+```console
+$ pip install "metrik-torch[torch]"
+$ metrik explore ./models/TinyLlama-1.1B --loader torch
+```
+
+The counterpart to `metrik-explorer`, which reads safetensors headers and `config.json` and
+never imports torch. This package walks the actual module tree, so module classes, parameter
+ownership, and **weight tying** are observed facts rather than claims from a config file.
+
+Tying in particular: an untied `lm_head` has exactly the same shape as the embedding it does
+not share, so shape is not evidence. Here it is established by storage identity.
+
+## What it still does not know
+
+Dataflow. `nn.Module` records containment, not what feeds what, so `edges` is `None` and
+`fidelity` is `structural`. Recovering real edges needs tracing, which breaks on the dynamic
+control flow most model implementations have. A consumer that needs edges checks `fidelity`
+and refuses rather than reading a containment tree as a dataflow graph.
+
+No forward pass, no device transfer, no weight value read — module structure and parameter
+metadata only.
+
+## Why torch is an optional extra
+
+`uv sync --all-packages --dev` stays torch-free so the CI matrix keeps running in about a
+minute. One dedicated job installs the extra and exercises extraction. See
+[ADR 007](../../docs/planning/adr/007-ml-substrate.md).

metrik_torch-0.1.0/README.md

@@ -0,0 +1,31 @@
+# metrik-torch
+
+Observed `ModelGraph` extraction from a real `nn.Module` tree.
+
+```console
+$ pip install "metrik-torch[torch]"
+$ metrik explore ./models/TinyLlama-1.1B --loader torch
+```
+
+The counterpart to `metrik-explorer`, which reads safetensors headers and `config.json` and
+never imports torch. This package walks the actual module tree, so module classes, parameter
+ownership, and **weight tying** are observed facts rather than claims from a config file.
+
+Tying in particular: an untied `lm_head` has exactly the same shape as the embedding it does
+not share, so shape is not evidence. Here it is established by storage identity.
+
+## What it still does not know
+
+Dataflow. `nn.Module` records containment, not what feeds what, so `edges` is `None` and
+`fidelity` is `structural`. Recovering real edges needs tracing, which breaks on the dynamic
+control flow most model implementations have. A consumer that needs edges checks `fidelity`
+and refuses rather than reading a containment tree as a dataflow graph.
+
+No forward pass, no device transfer, no weight value read — module structure and parameter
+metadata only.
+
+## Why torch is an optional extra
+
+`uv sync --all-packages --dev` stays torch-free so the CI matrix keeps running in about a
+minute. One dedicated job installs the extra and exercises extraction. See
+[ADR 007](../../docs/planning/adr/007-ml-substrate.md).

metrik_torch-0.1.0/pyproject.toml

@@ -0,0 +1,50 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "metrik-torch"
+version = "0.1.0"
+description = "Metrik torch loader — observed ModelGraph extraction from a real nn.Module tree."
+readme = "README.md"
+requires-python = ">=3.11"
+license = "Apache-2.0"
+authors = [{ name = "The Metrik Authors" }]
+classifiers = [
+    "Development Status :: 2 - Pre-Alpha",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+# metrik-sdk ONLY, same as every other first-party producer: no privileged access, enforced by
+# the `sdk-dogfooding` contract.
+#
+# torch is deliberately NOT here. It is an optional extra, so `uv sync --all-packages --dev`
+# stays torch-free and the seven-way CI matrix keeps running in about a minute. A single
+# dedicated job installs the extra and exercises extraction. The alternative -- torch as a hard
+# dependency -- would put a multi-hundred-megabyte download in front of every contributor and
+# every matrix job, to test one code path on one OS's worth of behaviour.
+dependencies = [
+    "metrik-sdk==0.1.0",
+]
+
+[project.optional-dependencies]
+# CPU is enough: this extracts structure, it never runs a forward pass. Anything that needs an
+# accelerator belongs in the profiler, behind the GPU tier.
+torch = ["torch>=2.2"]
+
+# Discovered through the registry like any third-party plugin, so `metrik explore --loader
+# torch` never imports this package by name.
+[project.entry-points."metrik.loaders"]
+torch = "metrik.torch:TorchLoader"
+
+[project.urls]
+Homepage = "https://github.com/Asmodeus14/Metrik"
+Repository = "https://github.com/Asmodeus14/Metrik"
+Issues = "https://github.com/Asmodeus14/Metrik/issues"
+
+[tool.uv.sources]
+metrik-sdk = { workspace = true }
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/metrik"]

metrik_torch-0.1.0/src/metrik/torch/__init__.py

@@ -0,0 +1,23 @@
+"""Metrik torch loader -- observed ModelGraph extraction from a real ``nn.Module`` tree.
+
+``torch`` is an optional extra. Importing this package without it raises immediately with a
+remediation rather than failing later inside extraction, because "No module named 'torch'"
+surfacing from three frames deep in a loader reads as a Metrik bug rather than a missing
+install.
+"""
+
+from __future__ import annotations
+
+from metrik.torch.extract import ROOT_NODE_ID, extract
+from metrik.torch.loader import INFO, TorchLoader, TorchLoaderParams
+from metrik.torch.roles import role_of_module, tensor_kind_of
+
+__all__ = [
+    "INFO",
+    "ROOT_NODE_ID",
+    "TorchLoader",
+    "TorchLoaderParams",
+    "extract",
+    "role_of_module",
+    "tensor_kind_of",
+]

metrik_torch-0.1.0/src/metrik/torch/extract.py

@@ -0,0 +1,273 @@
+"""Extracting a :class:`~metrik.sdk.ModelGraph` from a real ``nn.Module`` tree.
+
+``artifacts.md`` §6.2, ``method: torch_module_tree``. This is the *observed* counterpart to the
+Explorer's ``config_only`` path: the module tree is walked, so module classes, parameter
+ownership, and weight tying are facts about the object rather than claims from a config file.
+
+What it still does **not** know is dataflow. ``nn.Module`` records containment -- which module
+holds which -- and nothing about what feeds what. Recovering that needs tracing, which breaks
+on the dynamic control flow most model implementations have. So ``fidelity`` stays
+``structural`` and ``edges`` stays ``None``; a consumer that needs real edges checks the field
+and refuses rather than reading a containment tree as a dataflow graph.
+
+Nothing here runs a forward pass, moves a tensor to a device, or reads a weight value. The
+module tree and parameter *metadata* are all that is touched, so extraction is fast and cannot
+accidentally become a benchmark.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from metrik.sdk import (
+    IR_VERSION,
+    Block,
+    ExtractionInfo,
+    ExtractionMethod,
+    Fidelity,
+    ModelGraph,
+    Node,
+    NodeKind,
+    TensorDecl,
+    TensorKind,
+    Topology,
+)
+from metrik.torch.roles import role_of_module, tensor_kind_of
+
+if TYPE_CHECKING:  # pragma: no cover - torch is an optional extra
+    import torch.nn as nn
+
+__all__ = ["ROOT_NODE_ID", "extract"]
+
+#: The root module has no name in ``named_modules()``, and an empty ``node_id`` is not legal.
+#: Angle brackets cannot occur in a Python attribute path, so this cannot collide with a real
+#: module -- including the very common case of a root whose child is itself called ``model``.
+ROOT_NODE_ID = "<model>"
+
+#: torch dtype repr -> the safetensors spelling, so a graph from this loader and one from the
+#: header reader use one vocabulary. An unlisted dtype passes through as its torch name rather
+#: than being coerced into a wrong one.
+_DTYPE_NAMES = {
+    "torch.float32": "F32",
+    "torch.float16": "F16",
+    "torch.bfloat16": "BF16",
+    "torch.float64": "F64",
+    "torch.int8": "I8",
+    "torch.uint8": "U8",
+    "torch.int32": "I32",
+    "torch.int64": "I64",
+    "torch.bool": "BOOL",
+}
+
+
+def _dtype_name(dtype: Any) -> str:
+    return _DTYPE_NAMES.get(str(dtype), str(dtype))
+
+
+def _qualified(name: str) -> str:
+    return name or ROOT_NODE_ID
+
+
+def extract(
+    model: nn.Module,
+    *,
+    loader_name: str,
+    loader_version: str,
+    config: dict[str, Any] | None = None,
+    topology: Topology | None = None,
+) -> ModelGraph:
+    """Walk ``model`` and build the graph.
+
+    Args:
+        topology: Declared topology, usually read from ``config.json``. Merged *under* what
+            the module tree observes: an observed value always wins, because the config is a
+            claim and the tree is the object. ``tied_embeddings`` is the case that matters --
+            configs get it wrong, and here it is checked by identity.
+    """
+    nodes: list[Node] = []
+    tensors: list[TensorDecl] = []
+    unsupported: set[str] = set()
+
+    # Storage identity -> the first name that referred to it. Two parameters that are the same
+    # object are tied, and this is *observed*: no config key, no shape heuristic. Shape
+    # equality is not evidence -- an untied lm_head has the same shape as the embedding it does
+    # not share -- so identity is the only sound test, and torch gives it to us directly.
+    seen_storage: dict[int, str] = {}
+
+    for name, module in model.named_modules():
+        node_id = _qualified(name)
+        role = role_of_module(name, type(module).__name__)
+        if role.value == "other" and _has_own_parameters(module):
+            # Only report a miss for modules that actually own weights. Containers like
+            # `ModuleList` are legitimately unclassifiable and listing them would drown the
+            # signal that matters: a *parameterised* module we failed to understand.
+            unsupported.add(type(module).__name__)
+
+        own = _own_tensors(module)
+        for tensor_name, tensor, is_buffer in own:
+            full = f"{name}.{tensor_name}" if name else tensor_name
+            key = _storage_key(tensor)
+            tied_to = seen_storage.get(key)
+            if tied_to is None:
+                seen_storage[key] = full
+            tensors.append(
+                TensorDecl(
+                    name=full,
+                    shape=tuple(tensor.shape),
+                    dtype=_dtype_name(tensor.dtype),
+                    numel=tensor.numel(),
+                    # element_size() rather than a dtype table: it is right for every dtype
+                    # torch supports, including ones added after this was written.
+                    bytes_on_disk=tensor.numel() * tensor.element_size(),
+                    kind=TensorKind.BUFFER if is_buffer else tensor_kind_of(full, role),
+                    owner_node=node_id,
+                    tied_to=tied_to,
+                )
+            )
+
+        nodes.append(
+            Node(
+                node_id=node_id,
+                qualified_name=node_id,
+                kind=NodeKind.MODULE,
+                role=role,
+                module_class=type(module).__name__,
+                layer_index=_layer_index(name),
+                param_tensors=tuple(sorted(f"{name}.{n}" if name else n for n, _, _ in own)),
+            )
+        )
+
+    nodes.sort(key=lambda n: n.node_id)
+    tensors.sort(key=lambda t: t.name)
+    blocks = _blocks_of(model)
+
+    return ModelGraph(
+        ir_version=IR_VERSION,
+        nodes=tuple(nodes),
+        # Containment is not dataflow. See the module docstring.
+        edges=None,
+        tensors=tuple(tensors),
+        blocks=tuple(blocks),
+        topology=_merge_topology(topology, tensors, blocks, config),
+        extraction=ExtractionInfo(
+            method=ExtractionMethod.TORCH_MODULE_TREE,
+            fidelity=Fidelity.STRUCTURAL,
+            loader_name=loader_name,
+            loader_version=loader_version,
+            unsupported=tuple(sorted(unsupported)),
+        ),
+    )
+
+
+def _storage_key(tensor: Any) -> int:
+    """Identify the underlying storage.
+
+    ``data_ptr()`` rather than ``id(tensor)``: tying shares *storage*, and two Parameter
+    objects can wrap one allocation. A meta-device or otherwise unallocated tensor reports 0,
+    which would make every such tensor look tied to every other, so those fall back to object
+    identity and are simply never reported as tied.
+    """
+    pointer = int(tensor.data_ptr())
+    return pointer if pointer else -id(tensor)
+
+
+def _has_own_parameters(module: nn.Module) -> bool:
+    return any(True for _ in module.parameters(recurse=False))
+
+
+def _own_tensors(module: nn.Module) -> list[tuple[str, Any, bool]]:
+    """Parameters and buffers declared on this module, not on its children.
+
+    ``recurse=False`` is what makes ``owner_node`` meaningful: with recursion every ancestor
+    would claim every descendant's weights, and parameter counts per node would multiply by
+    tree depth.
+    """
+    out: list[tuple[str, Any, bool]] = [
+        (n, p, False) for n, p in module.named_parameters(recurse=False)
+    ]
+    out += [(n, b, True) for n, b in module.named_buffers(recurse=False)]
+    return out
+
+
+def _layer_index(name: str) -> int | None:
+    """The decoder layer this module sits in, if any.
+
+    Read from the qualified name, which is where torch puts ``ModuleList`` indices. Modules
+    outside the repeating stack get ``None`` rather than 0 -- collapsing them into layer 0
+    makes the first layer look anomalously large in every depth profile.
+    """
+    for part in name.split("."):
+        if part.isdigit():
+            return int(part)
+    return None
+
+
+def _blocks_of(model: nn.Module) -> list[Block]:
+    """Group the repeating stack.
+
+    Every ``ModuleList`` is a repeat group, keyed by its own qualified name. That is what lets
+    32 identical decoder layers be recognised as one structure rather than 32 unrelated ones,
+    which is what makes per-layer plans and visualizations comparable across model depths.
+    """
+    import torch.nn as nn_
+
+    blocks: list[Block] = []
+    for name, module in model.named_modules():
+        if not isinstance(module, nn_.ModuleList) or not len(module):
+            continue
+        for index, child in enumerate(module):
+            child_name = f"{name}.{index}" if name else str(index)
+            blocks.append(
+                Block(
+                    index=index,
+                    block_class=type(child).__name__,
+                    repeat_group_id=name or ROOT_NODE_ID,
+                    node_ids=tuple(
+                        sorted(
+                            f"{child_name}.{sub}" if sub else child_name
+                            for sub, _ in child.named_modules()
+                        )
+                    ),
+                )
+            )
+    return blocks
+
+
+def _merge_topology(
+    declared: Topology | None,
+    tensors: list[TensorDecl],
+    blocks: list[Block],
+    config: dict[str, Any] | None,
+) -> Topology:
+    """Combine what was declared with what was observed, observation winning.
+
+    Only two fields are genuinely observable from a module tree without tracing, and both are
+    worth overriding the config for:
+
+    ``n_layers`` -- counted from the repeating stack. A config can disagree with the weights it
+    ships, and the weights are what will be loaded.
+
+    ``tied_embeddings`` -- established by storage identity above. This is the field configs
+    most often get wrong, and getting it wrong changes the parameter count by the size of the
+    entire embedding matrix.
+    """
+    base = declared.model_dump() if declared is not None else {}
+
+    observed_layers = len({b.repeat_group_id for b in blocks}) and max(
+        (
+            len([b for b in blocks if b.repeat_group_id == g])
+            for g in {b.repeat_group_id for b in blocks}
+        ),
+        default=0,
+    )
+    if observed_layers:
+        base["n_layers"] = observed_layers
+
+    base["tied_embeddings"] = any(t.tied_to is not None for t in tensors)
+
+    if config and base.get("vocab_size") is None:
+        vocab = config.get("vocab_size")
+        if isinstance(vocab, int) and not isinstance(vocab, bool):
+            base["vocab_size"] = vocab
+
+    return Topology(**base)

metrik_torch-0.1.0/src/metrik/torch/loader.py

@@ -0,0 +1,211 @@
+"""The ``ModelLoader`` plugin the registry resolves for ``--loader torch``."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from metrik.sdk import (
+    DeterminismClass,
+    Emission,
+    Omission,
+    Params,
+    PartialInfo,
+    ProducerInfo,
+    ReasonCode,
+    UnsupportedError,
+    UserError,
+)
+
+__all__ = ["INFO", "TorchLoader", "TorchLoaderParams"]
+
+INFO = ProducerInfo(name="torch", version="0.0.1", api_version="1.0.0", kind="builtin")
+
+
+class TorchLoaderParams(Params):
+    """Everything that changes what is extracted.
+
+    ``trust_remote_code`` is here and defaults off because loading a model with custom code
+    executes the model author's Python. It is recorded in the artifact, so a graph extracted
+    with it on is distinguishable from one extracted without.
+    """
+
+    uri: str = ""
+    trust_remote_code: bool = False
+    #: Instantiate on the meta device: shapes and structure without allocating weights. On by
+    #: default because extraction never reads a weight value, so materialising gigabytes to
+    #: inspect the tree is pure cost. Turn it off only to inspect a model already in memory.
+    meta_device: bool = True
+
+
+def _require_torch() -> Any:
+    """Import torch, or fail with something a user can act on.
+
+    A bare ImportError three frames into extraction reads as a Metrik bug. This names the
+    install that fixes it.
+    """
+    try:
+        import torch
+    except ImportError as exc:  # pragma: no cover - exercised only without the extra
+        raise UnsupportedError(
+            "the torch loader needs PyTorch, which is not installed",
+            remediation=(
+                "Install it with `pip install 'metrik-torch[torch]'`, or use the header-only "
+                "loader with `metrik explore --loader explorer`."
+            ),
+        ) from exc
+    return torch
+
+
+class TorchLoader:
+    """Builds a ModelGraph by walking a real module tree."""
+
+    info = INFO
+    params_model = TorchLoaderParams
+
+    def can_load(self, uri: str) -> bool:
+        path = Path(uri).expanduser()
+        return (path / "config.json").is_file() if path.is_dir() else False
+
+    def graph_of_module(self, module: Any, *, config: dict[str, Any] | None = None) -> Any:
+        """Extract from an already-instantiated module.
+
+        The seam that makes this testable without transformers or a checkpoint: a plain
+        ``nn.Module`` built in a test exercises exactly the code a real model does.
+        """
+        from metrik.torch.extract import extract
+
+        return extract(
+            module,
+            loader_name=INFO.name,
+            loader_version=INFO.version,
+            config=config,
+            topology=_declared_topology(config or {}),
+        )
+
+    def load(self, uri: str, params: TorchLoaderParams | None = None) -> Any:
+        """Instantiate the model from a directory and extract its graph."""
+        params = params or TorchLoaderParams()
+        torch = _require_torch()
+        path = Path(uri).expanduser()
+        config_path = path / "config.json"
+        if not config_path.is_file():
+            raise UserError(
+                f"{uri} has no config.json, so there is nothing to instantiate",
+                remediation=(
+                    "Point at a model directory, or use `--loader explorer` to analyse "
+                    "weights without instantiating anything."
+                ),
+            )
+        config = json.loads(config_path.read_text(encoding="utf-8"))
+
+        try:
+            from transformers import AutoConfig, AutoModelForCausalLM
+        except ImportError as exc:
+            raise UnsupportedError(
+                "instantiating a model from a directory needs transformers",
+                remediation=(
+                    "Install it with `pip install transformers`, or use "
+                    "`metrik explore --loader explorer` for a header-only analysis."
+                ),
+            ) from exc
+
+        auto_config = AutoConfig.from_pretrained(
+            str(path), trust_remote_code=params.trust_remote_code
+        )
+        if params.meta_device:
+            # Structure only. No weights are allocated, which is what makes inspecting a 70B
+            # model on a laptop possible -- and is safe precisely because extraction never
+            # reads a value.
+            with torch.device("meta"):
+                # transformers ships no annotations for from_config, so these two calls are
+                # untyped in a strict context. Ignored narrowly rather than by relaxing the
+                # whole module: everything else here should stay checked.
+                model = AutoModelForCausalLM.from_config(  # type: ignore[no-untyped-call]
+                    auto_config, trust_remote_code=params.trust_remote_code
+                )
+        else:
+            model = AutoModelForCausalLM.from_config(  # type: ignore[no-untyped-call]
+                auto_config, trust_remote_code=params.trust_remote_code
+            )
+        return self.graph_of_module(model, config=config)
+
+    def provide_graph(self, spine: Any, uri: str, fingerprint_id: Any) -> Emission:
+        """The graph step, in the shape the explore chain expects.
+
+        Presence of this method is what marks a loader as able to supply a graph to someone
+        else's pipeline, so the CLI never has to know how any particular loader is built.
+
+        Resolves the fingerprint through the spine rather than being handed it, because the
+        spine records what a producer reads and checks it against what the producer declares.
+        Being passed the artifact directly would leave nothing observed and make that check
+        vacuous (``sdk-contracts.md`` §5).
+        """
+        fingerprint = spine.get(fingerprint_id)
+        graph = self.load(uri)
+        return self.emission(graph, fingerprint.envelope.content_hash)
+
+    def emission(self, graph: Any, fingerprint_hash: Any) -> Emission:
+        """Wrap an extracted graph as an artifact emission.
+
+        The omissions are not optional decoration. ``edges`` is null here, and a null field
+        with no recorded reason is indistinguishable from a bug (``artifacts.md`` §2.3) --
+        it also vanished from the CLI's "not determined" panel, so the gap became invisible
+        precisely where a reader would look for it. Anything a producer cannot determine is
+        named, with a reason, every time.
+        """
+        payload = graph.model_dump(mode="json", exclude_none=True)
+        payload["fingerprint"] = str(fingerprint_hash)
+
+        omissions = [
+            Omission(
+                field_path="edges",
+                reason_code=ReasonCode.MISSING_CAPABILITY,
+                detail=("nn.Module records containment, not dataflow, so no edges were observed"),
+                remediation=(
+                    "Dataflow needs tracing, which is not implemented and fails on the "
+                    "dynamic control flow most model implementations have."
+                ),
+            )
+        ]
+        if graph.extraction.unsupported:
+            named = ", ".join(graph.extraction.unsupported)
+            omissions.append(
+                Omission(
+                    field_path="nodes[].role",
+                    reason_code=ReasonCode.MISSING_CAPABILITY,
+                    detail=(
+                        f"{len(graph.extraction.unsupported)} parameterised module class(es) "
+                        f"were not classified: {named}"
+                    ),
+                    remediation=(
+                        "Those nodes are tagged `other`; check "
+                        "GraphSummary.classification_confidence before using this for a "
+                        "cross-architecture claim."
+                    ),
+                )
+            )
+
+        return Emission(
+            artifact_type="metrik.ModelGraph",
+            schema_version=1,
+            payload=payload,
+            determinism=DeterminismClass.DETERMINISTIC,
+            inputs={"fingerprint": fingerprint_hash},
+            partial=PartialInfo(complete=False, omissions=tuple(omissions)),
+        )
+
+
+def _declared_topology(config: dict[str, Any]) -> Any:
+    """Reuse the Explorer's config reader when it is installed.
+
+    Imported lazily and optionally rather than depended on: duplicating the alias tables would
+    guarantee they drift, but a hard dependency would make the torch loader unusable without
+    the header-only one. Absent it, the module tree still supplies what it observes.
+    """
+    try:
+        from metrik.explorer.topology import topology_of
+    except ImportError:  # pragma: no cover - explorer is installed in this workspace
+        return None
+    return topology_of(config)

metrik_torch-0.1.0/src/metrik/torch/metrik_plugin.json

@@ -0,0 +1,20 @@
+{
+  "torch": {
+    "name": "torch",
+    "kind": "model_loader",
+    "version": "0.0.1",
+    "metrik_api": ">=1.0,<2.0",
+    "summary": "Observed ModelGraph extraction from a real nn.Module tree. Structural fidelity: no dataflow edges.",
+    "homepage": "https://github.com/Asmodeus14/Metrik",
+    "license": "Apache-2.0",
+    "requires_device": [],
+    "requires_gradients": false,
+    "requires_calibration": false,
+    "requires_container": false,
+    "determinism": "deterministic",
+    "device_sensitive": false,
+    "writes_files": false,
+    "needs_network": false,
+    "executes_untrusted_code": true
+  }
+}

metrik_torch-0.1.0/src/metrik/torch/roles.py

@@ -0,0 +1,90 @@
+"""Mapping a module to its canonical :class:`~metrik.sdk.RoleTag`.
+
+Classified from the **qualified name**, not the module class. Almost every projection in a
+transformer is an ``nn.Linear``, so the class says nothing; the name says which one it is. The
+class is still recorded on the node, and is what identifies an unclassifiable *parameterised*
+module in ``ExtractionInfo.unsupported``.
+
+This is pattern matching on naming conventions, which is a heuristic and will eventually
+misclassify something. That is why nothing here guesses: a name that matches nothing is
+``OTHER``, its class is reported as unsupported, and the share of parameters landing in
+``OTHER`` is a confidence signal a consumer can act on (``artifacts.md`` §6.2).
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Final
+
+from metrik.sdk import RoleTag, TensorKind
+
+__all__ = ["role_of_module", "tensor_kind_of"]
+
+#: Matched in order, first hit wins. Order is load-bearing in two places: the fused patterns
+#: must precede the single-projection ones, or ``qkv_proj`` is read as ``q_proj``; and the
+#: expert patterns must precede the plain MLP ones, or every expert is filed as a dense MLP
+#: and a MoE model's parameter breakdown becomes meaningless.
+_PATTERNS: Final[tuple[tuple[re.Pattern[str], RoleTag], ...]] = tuple(
+    (re.compile(pattern), role)
+    for pattern, role in (
+        # MoE first.
+        (r"(^|\.)(gate|router)$", RoleTag.MOE_ROUTER),
+        (r"experts?\.\d+\.(w1|gate_proj)$", RoleTag.MOE_EXPERT_GATE),
+        (r"experts?\.\d+\.(w3|up_proj)$", RoleTag.MOE_EXPERT_UP),
+        (r"experts?\.\d+\.(w2|down_proj)$", RoleTag.MOE_EXPERT_DOWN),
+        (r"shared_expert", RoleTag.MOE_SHARED),
+        # Fused before single.
+        (r"(qkv_proj|c_attn|Wqkv|in_proj)$", RoleTag.ATTN_QKV_FUSED),
+        (r"gate_up_proj$", RoleTag.MLP_FUSED),
+        # Attention projections.
+        (r"(q_proj|query|wq)$", RoleTag.ATTN_Q),
+        (r"(k_proj|key|wk)$", RoleTag.ATTN_K),
+        (r"(v_proj|value|wv)$", RoleTag.ATTN_V),
+        (r"(o_proj|out_proj|dense|wo|c_proj)$", RoleTag.ATTN_O),
+        # MLP.
+        (r"gate_proj$", RoleTag.MLP_GATE),
+        (r"(up_proj|fc1|c_fc|w1)$", RoleTag.MLP_UP),
+        (r"(down_proj|fc2|w2)$", RoleTag.MLP_DOWN),
+        # Norms. `input_layernorm` guards the attention block, `post_attention_layernorm` the
+        # MLP -- naming that reads backwards until you notice it is named for what precedes it.
+        (r"(input_layernorm|ln_1|attention_norm|pre_attn_norm)$", RoleTag.ATTN_NORM),
+        (r"(post_attention_layernorm|ln_2|ffn_norm|pre_mlp_norm)$", RoleTag.MLP_NORM),
+        # Final norm: anchored to the top level, so a per-layer `model.layers.3.norm` cannot
+        # match it.
+        (r"^(model\.norm|transformer\.ln_f|norm|ln_f)$", RoleTag.NORM_FINAL),
+        # Embeddings and head.
+        (r"(embed_tokens|wte|tok_embeddings|word_embeddings)$", RoleTag.EMBED_TOKENS),
+        (r"(embed_positions|wpe|position_embeddings)$", RoleTag.EMBED_POS),
+        (r"(lm_head|output|score)$", RoleTag.LM_HEAD),
+    )
+)
+
+#: Module classes that hold parameters but carry no role of their own -- reporting them as
+#: unsupported would be noise, since there is nothing to classify.
+_CONTAINERS: Final[frozenset[str]] = frozenset({"ModuleList", "Sequential", "ModuleDict"})
+
+
+def role_of_module(qualified_name: str, module_class: str) -> RoleTag:
+    """Classify one module. Returns ``OTHER`` rather than guessing."""
+    if not qualified_name or module_class in _CONTAINERS:
+        return RoleTag.OTHER
+    for pattern, role in _PATTERNS:
+        if pattern.search(qualified_name):
+            return role
+    return RoleTag.OTHER
+
+
+def tensor_kind_of(tensor_name: str, role: RoleTag) -> TensorKind:
+    """Classify a parameter from its name and its owner's role.
+
+    Checked in this order deliberately: a norm's ``bias`` is a bias, not a norm scale, and a
+    norm's ``weight`` is a scale rather than a weight matrix. Getting that backwards puts
+    per-tensor statistics for a scalar scale in the same bucket as a projection matrix.
+    """
+    if tensor_name.endswith(".bias") or tensor_name == "bias":
+        return TensorKind.BIAS
+    if role in (RoleTag.ATTN_NORM, RoleTag.MLP_NORM, RoleTag.NORM_FINAL):
+        return TensorKind.NORM_SCALE
+    if role in (RoleTag.EMBED_TOKENS, RoleTag.EMBED_POS):
+        return TensorKind.EMBEDDING
+    return TensorKind.WEIGHT

metrik_torch-0.1.0/tests/test_extract.py

@@ -0,0 +1,332 @@
+"""Extraction against a hand-built module tree.
+
+No transformers, no checkpoint, no download. A plain ``nn.Module`` shaped like a decoder-only
+transformer exercises exactly the code a real model does, and keeps this runnable in the one
+CI job that installs torch.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+pytest.importorskip("torch", reason="metrik-torch[torch] is not installed")
+
+import torch
+from torch import nn
+
+from metrik.sdk import ExtractionMethod, Fidelity, RoleTag, TensorKind
+from metrik.torch import ROOT_NODE_ID, TorchLoader, extract
+
+pytestmark = pytest.mark.unit
+
+HIDDEN, HEADS, KV_HEADS, INTER, VOCAB, LAYERS = 64, 4, 2, 176, 512, 3
+
+
+class Attention(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        head_dim = HIDDEN // HEADS
+        self.q_proj = nn.Linear(HIDDEN, HEADS * head_dim, bias=False)
+        self.k_proj = nn.Linear(HIDDEN, KV_HEADS * head_dim, bias=False)
+        self.v_proj = nn.Linear(HIDDEN, KV_HEADS * head_dim, bias=False)
+        self.o_proj = nn.Linear(HEADS * head_dim, HIDDEN, bias=False)
+
+
+class Mlp(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(HIDDEN, INTER, bias=False)
+        self.up_proj = nn.Linear(HIDDEN, INTER, bias=False)
+        self.down_proj = nn.Linear(INTER, HIDDEN, bias=False)
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.self_attn = Attention()
+        self.mlp = Mlp()
+        self.input_layernorm = nn.LayerNorm(HIDDEN)
+        self.post_attention_layernorm = nn.LayerNorm(HIDDEN)
+
+
+class Inner(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.embed_tokens = nn.Embedding(VOCAB, HIDDEN)
+        self.layers = nn.ModuleList(DecoderLayer() for _ in range(LAYERS))
+        self.norm = nn.LayerNorm(HIDDEN)
+
+
+class TinyCausalLM(nn.Module):
+    def __init__(self, *, tie: bool = False) -> None:
+        super().__init__()
+        self.model = Inner()
+        self.lm_head = nn.Linear(HIDDEN, VOCAB, bias=False)
+        if tie:
+            # Exactly how transformers ties: the same Parameter object under two names.
+            self.lm_head.weight = self.model.embed_tokens.weight
+
+
+@pytest.fixture
+def graph() -> object:
+    return extract(TinyCausalLM(), loader_name="torch", loader_version="0.0.1")
+
+
+def _node(graph: object, node_id: str) -> object:
+    return next(n for n in graph.nodes if n.node_id == node_id)  # type: ignore[attr-defined]
+
+
+def _tensor(graph: object, name: str) -> object:
+    return next(t for t in graph.tensors if t.name == name)  # type: ignore[attr-defined]
+
+
+# -- structure ------------------------------------------------------------------------------
+
+
+def test_node_ids_are_qualified_names(graph: object) -> None:
+    ids = {n.node_id for n in graph.nodes}  # type: ignore[attr-defined]
+    assert "model.layers.1.self_attn.q_proj" in ids
+    assert "model.embed_tokens" in ids
+    # The root has no name in named_modules(), and an empty id is not legal.
+    assert ROOT_NODE_ID in ids
+
+
+def test_the_root_id_cannot_collide_with_a_real_module(graph: object) -> None:
+    """A root whose child is itself called `model` is the common case, so the sentinel uses
+    characters that cannot occur in a Python attribute path."""
+    assert "model" in {n.node_id for n in graph.nodes}  # type: ignore[attr-defined]
+    assert ROOT_NODE_ID != "model"
+
+
+def test_module_classes_are_recorded(graph: object) -> None:
+    """The thing the header-only loader fundamentally cannot know."""
+    assert _node(graph, "model.layers.0.self_attn.q_proj").module_class == "Linear"  # type: ignore[attr-defined]
+    assert _node(graph, "model.layers.0").module_class == "DecoderLayer"  # type: ignore[attr-defined]
+
+
+def test_parameters_are_owned_by_the_module_that_declares_them(graph: object) -> None:
+    """`recurse=False` is what makes owner_node meaningful: with recursion every ancestor
+    would claim every descendant's weights and per-node counts would multiply by depth."""
+    assert _tensor(graph, "model.layers.0.self_attn.q_proj.weight").owner_node == (  # type: ignore[attr-defined]
+        "model.layers.0.self_attn.q_proj"
+    )
+    assert _node(graph, "model.layers.0").param_tensors == ()  # type: ignore[attr-defined]
+
+
+def test_layer_index_is_none_outside_the_stack(graph: object) -> None:
+    """0 would make the first layer look anomalously large in every depth profile."""
+    assert _node(graph, "model.layers.2.mlp.up_proj").layer_index == 2  # type: ignore[attr-defined]
+    assert _node(graph, "model.embed_tokens").layer_index is None  # type: ignore[attr-defined]
+
+
+# -- roles ----------------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("node_id", "role"),
+    [
+        ("model.layers.0.self_attn.q_proj", RoleTag.ATTN_Q),
+        ("model.layers.0.self_attn.k_proj", RoleTag.ATTN_K),
+        ("model.layers.0.self_attn.v_proj", RoleTag.ATTN_V),
+        ("model.layers.0.self_attn.o_proj", RoleTag.ATTN_O),
+        ("model.layers.0.mlp.gate_proj", RoleTag.MLP_GATE),
+        ("model.layers.0.mlp.down_proj", RoleTag.MLP_DOWN),
+        ("model.layers.0.input_layernorm", RoleTag.ATTN_NORM),
+        ("model.layers.0.post_attention_layernorm", RoleTag.MLP_NORM),
+        ("model.embed_tokens", RoleTag.EMBED_TOKENS),
+        ("model.norm", RoleTag.NORM_FINAL),
+        ("lm_head", RoleTag.LM_HEAD),
+    ],
+)
+def test_roles(graph: object, node_id: str, role: RoleTag) -> None:
+    assert _node(graph, node_id).role is role  # type: ignore[attr-defined]
+
+
+def test_a_per_layer_norm_is_not_the_final_norm(graph: object) -> None:
+    """`model.norm` is anchored, so a hypothetical `model.layers.3.norm` cannot match it --
+    which would otherwise put a per-layer norm in the final-norm role."""
+    assert _node(graph, "model.norm").role is RoleTag.NORM_FINAL  # type: ignore[attr-defined]
+    assert _node(graph, "model.layers.0.input_layernorm").role is not RoleTag.NORM_FINAL  # type: ignore[attr-defined]
+
+
+def test_containers_are_not_reported_as_unsupported(graph: object) -> None:
+    """A ModuleList is legitimately unclassifiable. Listing it would drown the signal that
+    matters: a *parameterised* module we failed to understand."""
+    assert "ModuleList" not in graph.extraction.unsupported  # type: ignore[attr-defined]
+
+
+def test_an_unrecognised_parameterised_module_is_reported() -> None:
+    class Mystery(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weird_thing = nn.Linear(4, 4)
+
+    result = extract(Mystery(), loader_name="torch", loader_version="0.0.1")
+    assert "Linear" in result.extraction.unsupported
+    assert _node(result, "weird_thing").role is RoleTag.OTHER  # type: ignore[attr-defined]
+
+
+# -- weight tying, observed rather than declared ---------------------------------------------
+
+
+def test_tying_is_detected_by_storage_identity() -> None:
+    """The headline improvement over the header-only loader, which cannot see this at all:
+    an untied lm_head has exactly the same shape as the embedding it does not share."""
+    graph = extract(TinyCausalLM(tie=True), loader_name="torch", loader_version="0.0.1")
+    head = _tensor(graph, "lm_head.weight")
+    assert head.tied_to == "model.embed_tokens.weight"  # type: ignore[attr-defined]
+    assert graph.topology.tied_embeddings is True
+
+
+def test_an_untied_model_reports_no_tying(graph: object) -> None:
+    assert _tensor(graph, "lm_head.weight").tied_to is None  # type: ignore[attr-defined]
+    assert graph.topology.tied_embeddings is False  # type: ignore[attr-defined]
+
+
+def test_a_tied_weight_is_counted_once() -> None:
+    untied = extract(TinyCausalLM(), loader_name="torch", loader_version="0.0.1")
+    tied = extract(TinyCausalLM(tie=True), loader_name="torch", loader_version="0.0.1")
+    assert untied.param_count - tied.param_count == VOCAB * HIDDEN
+
+
+# -- blocks ---------------------------------------------------------------------------------
+
+
+def test_the_decoder_stack_is_one_repeat_group(graph: object) -> None:
+    """What lets N identical layers be recognised as one structure rather than N unrelated
+    ones, which is what makes per-layer plans comparable across model depths."""
+    groups = {b.repeat_group_id for b in graph.blocks}  # type: ignore[attr-defined]
+    assert groups == {"model.layers"}
+    assert len(graph.blocks) == LAYERS  # type: ignore[attr-defined]
+    assert {b.block_class for b in graph.blocks} == {"DecoderLayer"}  # type: ignore[attr-defined]
+
+
+def test_layer_count_is_observed_from_the_stack(graph: object) -> None:
+    assert graph.topology.n_layers == LAYERS  # type: ignore[attr-defined]
+
+
+# -- fidelity -------------------------------------------------------------------------------
+
+
+def test_extraction_is_structural_with_no_edges(graph: object) -> None:
+    """nn.Module records containment, not dataflow. Claiming edges here would be fabricating
+    them, and the IR rejects the combination outright."""
+    assert graph.extraction.method is ExtractionMethod.TORCH_MODULE_TREE  # type: ignore[attr-defined]
+    assert graph.extraction.fidelity is Fidelity.STRUCTURAL  # type: ignore[attr-defined]
+    assert graph.edges is None  # type: ignore[attr-defined]
+
+
+def test_tensor_kinds_distinguish_a_norm_scale_from_a_weight(graph: object) -> None:
+    assert _tensor(graph, "model.layers.0.input_layernorm.weight").kind is TensorKind.NORM_SCALE  # type: ignore[attr-defined]
+    assert _tensor(graph, "model.layers.0.input_layernorm.bias").kind is TensorKind.BIAS  # type: ignore[attr-defined]
+    assert _tensor(graph, "model.embed_tokens.weight").kind is TensorKind.EMBEDDING  # type: ignore[attr-defined]
+    assert _tensor(graph, "model.layers.0.mlp.up_proj.weight").kind is TensorKind.WEIGHT  # type: ignore[attr-defined]
+
+
+# -- determinism ----------------------------------------------------------------------------
+
+
+def test_two_extractions_of_one_model_are_identical() -> None:
+    """Nodes and tensors are sorted precisely so this holds: module traversal order is not
+    stable across transformers versions, and an unsorted graph would hash differently for an
+    unchanged model."""
+    model = TinyCausalLM()
+    first = extract(model, loader_name="torch", loader_version="0.0.1")
+    second = extract(model, loader_name="torch", loader_version="0.0.1")
+    assert first.model_dump_json() == second.model_dump_json()
+
+
+def test_extraction_never_reads_a_weight_value() -> None:
+    """Proved on the meta device, where the tensors have shapes but no storage at all. If
+    extraction touched a value this would raise; succeeding means only metadata was read --
+    which is what makes inspecting a 70B model on a laptop possible.
+    """
+    with torch.device("meta"):
+        model = TinyCausalLM(tie=True)
+    graph = extract(model, loader_name="torch", loader_version="0.0.1")
+    assert graph.param_count > 0
+    assert _node(graph, "model.layers.0.self_attn.q_proj").module_class == "Linear"  # type: ignore[attr-defined]
+
+
+def test_tying_survives_the_meta_device() -> None:
+    """The loader instantiates on `meta` by default, and every meta tensor reports
+    `data_ptr() == 0` -- so the obvious storage test would silently stop detecting tying on
+    the one path users actually take. It holds because tying shares the same Parameter object,
+    so identity still separates them, and the meta and real counts must agree exactly.
+    """
+    with torch.device("meta"):
+        meta = extract(TinyCausalLM(tie=True), loader_name="torch", loader_version="0.0.1")
+    real = extract(TinyCausalLM(tie=True), loader_name="torch", loader_version="0.0.1")
+
+    assert _tensor(meta, "lm_head.weight").tied_to == "model.embed_tokens.weight"  # type: ignore[attr-defined]
+    assert meta.topology.tied_embeddings is True
+    assert meta.param_count == real.param_count
+
+
+def test_meta_tensors_are_not_all_reported_as_tied_to_each_other() -> None:
+    """Every meta tensor reports data_ptr() == 0. Keying tying on that alone would make the
+    entire model look like one shared weight, and the parameter count collapse to a single
+    tensor -- a wrong number that looks like a very efficient model.
+    """
+    with torch.device("meta"):
+        model = TinyCausalLM()
+    graph = extract(model, loader_name="torch", loader_version="0.0.1")
+    tied = [t for t in graph.tensors if t.tied_to is not None]
+    assert tied == []
+
+
+# -- the loader contract ----------------------------------------------------------------------
+
+
+def test_the_loader_extracts_from_an_instantiated_module() -> None:
+    graph = TorchLoader().graph_of_module(TinyCausalLM(), config={"vocab_size": VOCAB})
+    assert graph.topology.vocab_size == VOCAB
+    assert graph.extraction.loader_name == "torch"
+
+
+def test_can_load_requires_a_directory_with_a_config(tmp_path: Path) -> None:
+    loader = TorchLoader()
+    assert not loader.can_load(str(tmp_path))
+    (tmp_path / "config.json").write_text("{}", encoding="utf-8")
+    assert loader.can_load(str(tmp_path))
+
+
+# -- honesty ---------------------------------------------------------------------------------
+
+
+def test_the_emission_declares_its_missing_edges() -> None:
+    """A null field with no recorded reason is indistinguishable from a bug.
+
+    Caught by running it: the graph had `edges: None` and no omission, so it silently dropped
+    out of the CLI's "not determined" panel -- disappearing from exactly the place a reader
+    looks to find out what was not determined.
+    """
+    from metrik.core.hashing import Digest
+
+    loader = TorchLoader()
+    graph = loader.graph_of_module(TinyCausalLM())
+    emission = loader.emission(graph, Digest.parse("b3:" + "11" * 32))
+
+    assert emission.partial is not None
+    assert not emission.partial.complete
+    paths = {o.field_path for o in emission.partial.omissions}
+    assert "edges" in paths
+
+
+def test_unclassified_modules_are_surfaced_as_an_omission() -> None:
+    """A graph where a parameterised module went unrecognised should say so, because the
+    share of parameters tagged `other` is what a consumer checks before making a
+    cross-architecture claim."""
+    from metrik.core.hashing import Digest
+
+    class Mystery(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weird_thing = nn.Linear(4, 4)
+
+    loader = TorchLoader()
+    emission = loader.emission(loader.graph_of_module(Mystery()), Digest.parse("b3:" + "22" * 32))
+    assert emission.partial is not None
+    assert "nodes[].role" in {o.field_path for o in emission.partial.omissions}