Skip to content

Fix shared-parameter checkpoints and check shapes on load - #97

Open
cemde wants to merge 1 commit into
liukidar:mainfrom
cemde:fix/73-serialisation
Open

Fix shared-parameter checkpoints and check shapes on load#97
cemde wants to merge 1 commit into
liukidar:mainfrom
cemde:fix/73-serialisation

Conversation

@cemde

@cemde cemde commented Aug 9, 2026

Copy link
Copy Markdown
Collaborator

Bug

Two defects in _serialisation.py.

save_params deduplicates shared parameters by writing the value under the first path it meets and a literal None under every later one. np.savez_compressed stores that None as a 0-d dtype=object array. load_params then calls np.load, which refuses object arrays unless allow_pickle=True.

Separately, load_params assigned whatever the file held without comparing shapes.

Impact

Every checkpoint of a model using pxnn.shared is unrecoverable:

ValueError: Object arrays cannot be loaded when allow_pickle=False

The loader's is not None guard was dead code, since np.load yields a 0-d object array and never None, so even with allow_pickle=True the shared parameter would have been assigned garbage.

Without a shape check, a Linear(5, 7) checkpoint loaded into a Linear(2, 3) leaves the weight at (7, 5), surfacing much later as an inscrutable broadcasting error inside a training step.

Fix

The sentinel the writer and reader had to agree on turned out to be no key at all. save_params drops the else branch, so duplicates are simply not written; load_params mirrors it with the same _cache() idiom, skipping the same paths. No object array is ever created, so nothing needs allow_pickle.

Plus a shape comparison raising ValueError before assignment.

Compatibility

Old checkpoints still load, including shared ones that previously could not. np.load on an .npz is lazy and the loader no longer visits the duplicate paths, so the stale dtype=object entries are never read.

The shape check runs only when the target parameter already holds a value. A freshly constructed Vode has h set to None, and model = Model(); load_params(model, ...) is the documented workflow, so an unconditional _param.shape raised AttributeError on the normal path. Found in review; regression test added, since every sibling test seeded its target first and left that path uncovered.

Closes #73

Duplicate references were written as a literal None, which numpy stores as
a dtype=object array that np.load refuses without allow_pickle. A shared
parameter is now saved once, under its first path, and the loader skips the
paths it has already restored. Loading a value whose shape does not match
its target parameter now raises instead of corrupting the model.
@cemde
cemde force-pushed the fix/73-serialisation branch from 26eaf31 to 603cbb5 Compare August 9, 2026 18:23
@cemde
cemde requested a review from liukidar August 9, 2026 18:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Checkpoints: shared parameters cannot be loaded, and shapes are not checked

1 participant