Add public model graph and editing API

[Vi-Sri]

Summary

• add a versioned, JSON-serializable model graph and per-layer/per-neuron inspection API
• add validated dependency-aware helpers for add, prune, freeze, unfreeze, reset, and perturb operations
• make freeze/unfreeze idempotent at the public API boundary
• document the model graph ontology, dependency types, and editing workflow
• add regression coverage for graph serialization, live structural edits, frozen state, and invalid targets

Why

Issue #267 identifies the existing modelling implementation as difficult to discover and use. The dependency engine and neuron operations already exist, but callers must understand internal layer ids, enums, and operate() semantics. This PR exposes the first vertical API slice requested by the issue while preserving the existing Torch FX/ONNX dependency machinery.

Issue #267 remains open after this PR. The deeper index-mapping, hardcoding, monitoring, lineage, and broad architecture coverage tracked by #6, #7, and #8 remain follow-up work.

Developer impact

A wrapped model created with compute_dependencies=True can now describe its editable graph through get_model_graph() / get_layer_info() and perform validated edits through named methods. Structural graph inspection stays lightweight by default; per-neuron records are opt-in for whole-graph responses.

Validation

• 38 passed, 7 skipped across focused model, interface, constraint, and trainer-service tests
• Ruff CI error rules, compileall, and git diff --check pass
• Full suite reached 949 passed, 131 skipped; remaining local failures/errors require unavailable MNIST downloads, the separately built/gitignored Studio bundle, or restricted local port/multiprocessing capabilities

[Copilot]

Pull request overview

Adds a first-class, public “model graph + editing” API on top of the existing dependency machinery, so callers can discover editable layers by stable-in-process IDs, inspect layer/neuron metadata, and perform validated architecture/weight edits without relying on internal enums or operate() details.

Changes:

• Added JSON-serializable model graph and per-layer inspection helpers (get_model_graph(), get_layer_info()), including dependency edge export and optional per-neuron details.
• Added validated public editing helpers (add_neurons, prune_neurons, freeze_neurons, unfreeze_neurons, reset_neurons, perturb_neurons) with idempotent freeze/unfreeze behavior at the API boundary.
• Added documentation and regression tests for graph structure/serialization, live edits, idempotency, and invalid target validation.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
weightslab/models/model_with_ops.py	Introduces the public model graph + editing API surface and validation/idempotency logic on top of the existing dependency manager and operate() implementation.
tests/model/test_model_graph_api.py	Adds focused regression coverage for graph serialization, layer info shape, idempotent freeze behavior, live structural edits, and input validation.
docs/model_interaction.rst	Documents the new graph ontology and editing workflow, including dependency types and modifier usage examples.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Suppressed comments (1)

weightslab/models/model_with_ops.py:176

• _build_layer_info() always materializes neuron_lrs for every output neuron to compute frozen_neuron_count, even when include_neurons=False (the default for get_model_graph). This makes the “cheap by default” structural graph potentially O(total_neurons) in time and memory for large models, despite only needing a frozen count in that mode. You can avoid the full scan by only building neuron_lrs when include_neurons=True, and otherwise deriving frozen indices/count from the sparse neuron_2_lr['weight'] overrides (defaults are 1.0).

        lr_overrides = getattr(layer, "neuron_2_lr", {})
        weight_lrs = lr_overrides.get("weight", {}) if lr_overrides else {}
        neuron_count = output_neurons or 0
        neuron_lrs = [float(weight_lrs.get(index, 1.0)) for index in range(neuron_count)]
        frozen_neurons = {index for index, lr in enumerate(neuron_lrs) if lr == 0.0}

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (1)

weightslab/models/model_with_ops.py:176

• get_model_graph(include_neurons=False) is documented as a cheap structural inspection, but _build_layer_info() still iterates over every output neuron to build neuron_lrs/frozen_neurons even when include_neurons is false. For large layers this makes the default graph response O(total_neurons) and can also grow the underlying defaultdict access patterns unnecessarily. Compute frozen counts from the override map when neurons are excluded, and only build the per-neuron list when include_neurons=True.

        lr_overrides = getattr(layer, "neuron_2_lr", {})
        weight_lrs = lr_overrides.get("weight", {}) if lr_overrides else {}
        neuron_count = output_neurons or 0
        neuron_lrs = [float(weight_lrs.get(index, 1.0)) for index in range(neuron_count)]
        frozen_neurons = {index for index, lr in enumerate(neuron_lrs) if lr == 0.0}

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (2)

weightslab/examples/PyTorch/wl-model-editing/vit-model-editing/README.md:8

• The README uses the term "ViT B-12", but the rest of the document (and torchvision naming) refers to ViT-B/16 (12 encoder blocks, 16x16 patches). This is likely to confuse readers about which model variant the experiment targets.

"ViT B-12" is interpreted here as **ViT-Base with 12 transformer encoder
blocks**. The torchvision implementation is `vit_b_16`: Base width/depth with
16x16 input patches.

weightslab/models/model_with_ops.py:266

• _validated_neuron_indices accepts bool values because bool is a subclass of int (e.g., [True] passes validation). This can lead to surprising behavior (freezing/pruning neuron 1) and undermines the API's input validation.

        if not all(isinstance(index, int) for index in indices):
            raise TypeError("neuron_indices must contain integers only.")