Add public model graph and editing API - #287
Draft
Vi-Sri wants to merge 2 commits into
Draft
Conversation
Contributor
There was a problem hiding this comment.
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.
Contributor
There was a problem hiding this comment.
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 materializesneuron_lrsfor every output neuron to computefrozen_neuron_count, even wheninclude_neurons=False(the default forget_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 buildingneuron_lrswheninclude_neurons=True, and otherwise deriving frozen indices/count from the sparseneuron_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}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
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=Truecan now describe its editable graph throughget_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 skippedacross focused model, interface, constraint, and trainer-service testscompileall, andgit diff --checkpass949 passed, 131 skipped; remaining local failures/errors require unavailable MNIST downloads, the separately built/gitignored Studio bundle, or restricted local port/multiprocessing capabilities