Serialization
4 functions — 4 documented.
save¶
Write named arrays to a vkML checkpoint.
The file is a ZIP archive containing vkml.json — format identifier, version, key list and user metadata — and one tensors/<key>.npy per entry.
That layout is a security decision before an engineering one. The well-known failure in this field is a model format that deserialises into code execution: Python's pickle reconstructs objects by calling what the stream names, so loading a checkpoint from anywhere but your own disk runs a program someone else wrote. A vkML checkpoint has no mechanism for it — every member is either a .npy array read with allow_pickle=False or one JSON document, and neither can name a callable.
Parsing the array bytes is delegated to NumPy's own reader rather than hand-written, because a hand-rolled binary parser is exactly where a memory-safety bug would go.
Parameters
- path (str | Path)
- Destination file.
- tensors (Mapping[str, numpy.ndarray])
- Arrays to store, by name.
- metadata (Mapping[str, Any] = None) optional
- JSON-representable extras — epoch, score, architecture name.
- compress (bool = False) optional
- Deflate the payload. Off by default.
The path comes first, and the payload is NumPy arrays rather than tensors.
Example
>>> import tempfile, os
>>> path = os.path.join(tempfile.mkdtemp(), "ckpt.vkml")
>>> w = vkml.tensor(np.zeros((4, 4), dtype=np.float32))
>>> vkml.save(path, {"w": w.numpy()}, metadata={"epoch": 3})
>>> ck = vkml.load(path)
>>> ck.metadata["epoch"]
3Implementation
| CPU kernel | composed from other operators |
|---|---|
| Vulkan shader | composed, or dispatched through a shared kernel |
| Gradient rule | none — backward through it raises |
| Tests (≥22) | test_data_and_serialize.py test_platform_portability.py |
See also load, save_module, load_module
load¶
Read a vkML checkpoint and return its arrays and metadata.
Returns a Checkpoint with .tensors, .metadata and .version, kept as separate fields so a metadata key can never collide with a tensor name.
Parameters
- path (str | Path)
- Checkpoint to read.
- max_expansion_ratio (float = 100.0) optional
- Reject the file if its members expand by more than this factor.
Returns
A Checkpoint.
Known false positive: a pruned or sparse model stored densely is mostly zeros and can exceed the ratio legitimately. The error names the file, the ratio and the argument that raises it.
Decompression bombs are rejected by expansion ratio, not by absolute size, and the reason is measured rather than assumed: every checkpoint in this repository expands 1.00× (the default is stored, not deflated), real float32 weights asked to compress reach 1.1× — weights are high-entropy and barely compress — and an all-zeros bomb reaches 1017×. Three orders of magnitude with nothing in between.
A byte cap has no non-arbitrary value: set low it breaks a legitimate large model, set high it stops nothing, and a 28 GB checkpoint and a 200 KB bomb are both things someone might load. The ratio makes an attacker's cost scale with the damage — forcing an N-byte allocation requires shipping N/100 bytes. It does not bound memory absolutely and is not meant to.
Implementation
| CPU kernel | composed from other operators |
|---|---|
| Vulkan shader | composed, or dispatched through a shared kernel |
| Gradient rule | none — backward through it raises |
| Tests (≥24) | test_data_and_serialize.py test_platform_portability.py |
See also save, load_module
save_module¶
Write a module's state dict to a checkpoint.
Parameters
- path (str | Path)
- Destination file.
- module (Module)
- The module whose state to save.
- metadata (Mapping[str, Any] = None) optional
- JSON-representable extras.
- compress (bool = False) optional
- Deflate the payload.
Implementation
| CPU kernel | composed from other operators |
|---|---|
| Vulkan shader | composed, or dispatched through a shared kernel |
| Gradient rule | none — backward through it raises |
| Tests (≥4) | test_data_and_serialize.py |
See also load_module, save
load_module¶
Load a checkpoint into an existing module and return it for its metadata.
Parameters are restored in place, keeping each entry's device, dtype and requires_grad. A module already moved to a GPU stays there.
The state dict must match exactly — a missing or unexpected key raises rather than being ignored, so a checkpoint from a different architecture fails at the load instead of producing a partly-initialised model.
Parameters
- path (str | Path)
- Checkpoint to read.
- module (Module)
- Module to load into.
Returns
The Checkpoint, for its metadata.
The metadata arrives after the state dict is installed, so a check written against it cannot guard the load. To decide whether to load at all, call load first, inspect, then load_state_dict yourself.
Example
>>> import tempfile, os
>>> path = os.path.join(tempfile.mkdtemp(), "m.vkml")
>>> dev = vkml.device("vulkan:0")
>>> model = vkml.nn.Linear(16, 8).to(dev)
>>> vkml.save_module(path, model)
>>> ck = vkml.load_module(path, model)
>>> next(iter(model.named_parameters()))[1].device # unchanged by the load
device('vulkan:0')Implementation
| CPU kernel | composed from other operators |
|---|---|
| Vulkan shader | composed, or dispatched through a shared kernel |
| Gradient rule | none — backward through it raises |
| Tests (≥4) | test_data_and_serialize.py |
See also save_module, load, save