Add public model graph and editing API

docs/model_interaction.rst

@@ -36,9 +36,59 @@ Minimal example
        log=True,
    )
 
+Model graph and editing API
+---------------------------
+
+Model editing needs dependency information so a change to one layer can be
+propagated to connected layers. Enable it when wrapping the model:
+
+.. code-block:: python
+
+   model = wl.watch_or_edit(
+       my_model,
+       flag="model",
+       dummy_input=example_batch,
+       compute_dependencies=True,
+   )
+
+   graph = model.get_model_graph()
+   first_layer = model.get_layer_info(graph["layers"][0]["id"])
+
+``get_model_graph()`` returns JSON-serializable primitives with a schema
+version, model metadata, layers, and directed dependencies. Each dependency is
+one of:
+
+- ``SAME``: both layers expose the same neuron/channel dimension, such as a
+  convolution followed by batch normalization.
+- ``INCOMING``: the source output feeds the destination input, such as one
+  linear layer feeding another.
+- ``REC``: a recursive or skip-connection relationship that must be kept in
+  sync during structural edits.
+
+The structural graph omits per-neuron records by default to keep inspection
+cheap. Use ``get_model_graph(include_neurons=True)`` or
+``get_layer_info(layer_id)`` for learning-rate and frozen-state details.
+
+The model exposes dependency-aware modifiers:
+
+.. code-block:: python
+
+   model.freeze_neurons(layer_id=0, neuron_indices=[0, 2])
+   model.unfreeze_neurons(layer_id=0, neuron_indices=[0])
+   model.reset_neurons(layer_id=0, neuron_indices=[1])
+   model.perturb_neurons(layer_id=0, neuron_indices=[2], ratio=0.1)
+   model.add_neurons(layer_id=0, count=2)
+   model.prune_neurons(layer_id=0, neuron_indices=[3])
+
+Omitting ``neuron_indices`` freezes, unfreezes, resets, or perturbs the whole
+layer. Layer ids are stable only for the lifetime of the wrapped model; resolve
+them from the graph instead of persisting them as checkpoint identifiers.
+
 Best practices
 --------------
 
 - Use explicit names for losses/metrics to keep logs readable.
 - Prefer ``per_sample=True`` for losses when you need hard-example analysis.
 - Keep model/device arguments explicit to avoid ambiguity in multi-device setups.
+- Pause training before structural edits so model and optimizer updates happen
+  at a safe boundary.

tests/model/test_model_graph_api.py

@@ -0,0 +1,117 @@
+import json
+import unittest
+
+import torch
+import torch.nn as nn
+
+from weightslab.backend import ledgers
+from weightslab.backend.model_interface import ModelInterface
+
+
+class TestModelGraphApi(unittest.TestCase):
+    def tearDown(self):
+        ledgers.clear_all()
+
+    @staticmethod
+    def _wrapped_model():
+        model = nn.Sequential(
+            nn.Linear(4, 3),
+            nn.ReLU(),
+            nn.Linear(3, 2),
+        )
+        wrapped = ModelInterface(
+            model,
+            dummy_input=torch.randn(1, 4),
+            compute_dependencies=True,
+            register=False,
+            skip_previous_auto_load=True,
+        )
+        # Architecture-change hooks update registered optimizers. These focused
+        # API tests do not register one, so avoid unrelated ledger work.
+        wrapped._architecture_change_hook_fns = []
+        return wrapped
+
+    def test_get_model_graph_returns_serializable_structure(self):
+        model = self._wrapped_model()
+
+        graph = model.get_model_graph()
+
+        self.assertEqual(graph["schema_version"], 1)
+        self.assertEqual([layer["name"] for layer in graph["layers"]], ["0", "1", "2"])
+        self.assertEqual(
+            graph["dependencies"],
+            [
+                {"source_layer_id": 0, "target_layer_id": 1, "type": "SAME"},
+                {"source_layer_id": 1, "target_layer_id": 2, "type": "INCOMING"},
+            ],
+        )
+        self.assertNotIn("neurons", graph["layers"][0])
+        json.dumps(graph)
+
+    def test_layer_info_and_modifiers_update_the_live_model(self):
+        model = self._wrapped_model()
+
+        layer = model.get_layer_info(0)
+        self.assertEqual(layer["type"], "Linear")
+        self.assertEqual(layer["input_neurons"], 4)
+        self.assertEqual(layer["output_neurons"], 3)
+        self.assertEqual(len(layer["neurons"]), 3)
+
+        model.freeze_neurons(0, [0])
+        self.assertTrue(model.get_layer_info(0)["neurons"][0]["frozen"])
+
+        # Freezing twice is idempotent rather than toggling the state.
+        model.freeze_neurons(0, [0])
+        self.assertTrue(model.get_layer_info(0)["neurons"][0]["frozen"])
+        self.assertEqual(model.get_layer_info(0)["operation_counts"]["FREEZE"], 1)
+
+        model.unfreeze_neurons(0, [0])
+        self.assertFalse(model.get_layer_info(0)["neurons"][0]["frozen"])
+
+        model.add_neurons(0, count=2)
+        self.assertEqual(model.get_layer_info(0)["output_neurons"], 5)
+        self.assertEqual(model.get_layer_info(2)["input_neurons"], 5)
+        self.assertEqual(tuple(model(torch.randn(1, 4)).shape), (1, 2))
+
+    def test_modelling_api_rejects_ambiguous_or_invalid_targets(self):
+        model = self._wrapped_model()
+
+        with self.assertRaisesRegex(ValueError, "Unknown layer id"):
+            model.get_layer_info(99)
+        with self.assertRaisesRegex(ValueError, "cannot be empty"):
+            model.prune_neurons(0, [])
+        with self.assertRaisesRegex(ValueError, "outside layer 0"):
+            model.reset_neurons(0, [3])
+        with self.assertRaisesRegex(ValueError, "no learnable weights"):
+            model.freeze_neurons(1, [0])
+        with self.assertRaisesRegex(ValueError, "strictly between 0 and 1"):
+            model.perturb_neurons(0, [0], ratio=1.0)
+
+    def test_freeze_and_unfreeze_reject_untracked_weight(self):
+        raw_model = nn.Sequential(
+            nn.Linear(4, 3),
+            nn.ReLU(),
+            nn.Linear(3, 2),
+        )
+        raw_model[0].weight.requires_grad_(False)
+        model = ModelInterface(
+            raw_model,
+            dummy_input=torch.randn(1, 4),
+            compute_dependencies=True,
+            register=False,
+            skip_previous_auto_load=True,
+        )
+        model._architecture_change_hook_fns = []
+
+        for operation in (model.freeze_neurons, model.unfreeze_neurons):
+            with self.subTest(operation=operation.__name__):
+                with self.assertRaisesRegex(
+                    ValueError, "trainable, per-neuron tracked weight"
+                ):
+                    operation(0, [0])
+
+        self.assertEqual(model.get_layer_info(0)["operation_counts"]["FREEZE"], 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

weightslab/examples/PyTorch/wl-model-editing/README.md

@@ -0,0 +1,13 @@
+# Model-editing experiments
+
+This directory contains architecture-specific experiments that validate
+WeightsLab model editing against live forward, backward, optimizer, and resumed
+training behavior.
+
+Available experiments:
+
+- [`vit-model-editing`](vit-model-editing/README.md): ViT-B/16 full-model
+  compatibility probe and supported editable-head control experiment.
+
+Add future model families as sibling directories so their compatibility limits
+and behavioral assertions remain independently runnable.

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

@@ -0,0 +1,119 @@
+# ViT model-editing experiment
+
+This experiment tests WeightsLab model editing instead of assuming that an
+operation succeeded because the API returned without raising.
+
+"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.
+
+## What is tested
+
+There are two separate scripts because they answer different questions.
+
+`probe_full_vit.py` is the compatibility gate for wrapping and editing the
+complete ViT-B/16 model. It verifies:
+
+1. Ordinary ViT forward/backward training works.
+2. WeightsLab can build its editable dependency graph.
+3. A classifier-head freeze can be applied.
+4. Training can resume after the edit.
+
+The current implementation is expected to stop at step 2. Torch FX traces the
+model, but WeightsLab dependency mapping reaches `MultiheadAttention`, which
+does not expose the neuron-operation interface expected by
+`generate_index_maps`. The JSON report records the exact exception. Run without
+`--allow-unsupported` when using this as a CI gate.
+
+`run_head_experiment.py` is the working control experiment. It:
+
+1. Builds a real 12-block ViT-B/16 backbone.
+2. Generates deterministic patterned images without a network download.
+3. Caches frozen CLS embeddings from the backbone.
+4. Trains a WeightsLab-wrapped MLP classification head.
+5. Adds hidden neurons and checks that the downstream input shape changes.
+6. Checks that the optimizer is rebuilt with the new parameters.
+7. Freezes a neuron and asserts that its gradient is zero.
+8. Unfreezes it and asserts that its gradient becomes non-zero.
+9. Resumes training and writes all graph, loss, and assertion results to JSON.
+
+This control is intentionally limited to the head. Editing the 768-dimensional
+transformer representation would require synchronized changes across attention
+projections, residual paths, positional embeddings, and LayerNorm parameters;
+the current dependency engine does not model those relationships reliably.
+
+## Setup
+
+From the repository root:
+
+```bash
+bash weightslab/examples/PyTorch/wl-model-editing/vit-model-editing/setup.sh
+```
+
+The default environment is `.venv-vit-edit`. Override it with
+`VIT_EDIT_VENV=/absolute/path` if needed.
+
+## Run both experiments
+
+```bash
+bash weightslab/examples/PyTorch/wl-model-editing/vit-model-editing/run_all.sh
+```
+
+Reports are written under `outputs/vit_model_editing/`, which is gitignored:
+
+- `full_vit_compatibility_report.json`
+- `vit_head_editing_report.json`
+
+The combined runner allows the known full-ViT incompatibility so the supported
+head experiment still runs. For a strict full-model gate:
+
+```bash
+.venv-vit-edit/bin/python \
+  weightslab/examples/PyTorch/wl-model-editing/vit-model-editing/probe_full_vit.py
+```
+
+That command exits non-zero until complete ViT dependency mapping works.
+
+## Useful variants
+
+Fast local smoke run (default, no downloads):
+
+```bash
+.venv-vit-edit/bin/python \
+  weightslab/examples/PyTorch/wl-model-editing/vit-model-editing/run_head_experiment.py \
+  --image-size 32 --train-samples 16 --eval-samples 8
+```
+
+Standard 224x224 input geometry with randomly initialized weights:
+
+```bash
+.venv-vit-edit/bin/python \
+  weightslab/examples/PyTorch/wl-model-editing/vit-model-editing/run_head_experiment.py \
+  --image-size 224
+```
+
+Pretrained ImageNet backbone (downloads torchvision weights):
+
+```bash
+.venv-vit-edit/bin/python \
+  weightslab/examples/PyTorch/wl-model-editing/vit-model-editing/run_head_experiment.py \
+  --image-size 224 --pretrained
+```
+
+CPU and CUDA are supported. MPS is intentionally excluded because the current
+`ModelInterface` device normalization maps non-CUDA devices back to CPU.
+
+## Reading a result
+
+A passing head report must have:
+
+- `status: "passed"`
+- `checks.add_propagated: true`
+- `checks.optimizer_rebuilt: true`
+- `checks.frozen_gradient_norm: 0.0`
+- `checks.unfrozen_gradient_norm > 0`
+- `checks.forward_after_edit: true`
+- `checks.training_resumed: true`
+
+These are behavioral checks against the live model and optimizer, not only API
+shape checks.