NURL registrynurl-lang.org →packages/torchpt/src/torchpt.nu — read PyTorch .pt / .pth checkpoints in pure NURL.
A modern torch.save file is a ZIP archive:
<name>/data.pkl the object graph, pickled <name>/data/<key> one flat storage per key, raw little-endian <name>/version etc.
Tensors in the pickle are _rebuild_tensor_v2(storage, offset, size, stride, …) calls whose storage is a persistent id naming one of the data/<key> members. So reading a checkpoint is: parse the zip, parse the pickle as DATA (never executing it — see src/pickle.nu), then map each tensor onto a byte range of the mapping.
mmap-backed and lazy, like safetensor and gguf: only the pickle is read up front, so listing the tensors of a 4.6 GB checkpoint costs megabytes, not gigabytes. Tensor bytes are addressed straight out of the mapping.
Entries are ZIP-stored (torch never deflates tensor data), which is what makes in-place addressing possible; a deflated storage is reported rather than silently mis-read.
( pt_open path ) → !Pt String ( pt_close p ) → v early release (optional) ( pt_none ) / ( pt_is_open p ) → Pt / b an empty slot, and the test for one ( pt_n_tensors p ) → i ( pt_name p idx ) → s BORROWED ( pt_find p name ) → i -1 when absent ( pt_dtype p idx ) → i PKS_ ( pt_ndim p idx ) → i ( pt_dim p idx j ) → i ( pt_nelems p idx ) → i ( pt_nbytes p idx ) → i ( pt_is_contiguous p idx ) → b ( pt_tensor_ptr p idx ) → u into the mapping ( pt_dequant p idx ) → !( Vec u ) String f32 LE bytes ( pt_dequant_range p idx first count ) → !( Vec u ) String ( pt_read_f64 p idx first count dst ) → b into an f64 buffer
Names are the dotted path from the pickle root: a plain state_dict gives blocks.0.attn.qkv.weight, while a training checkpoint saved as {"model": sd, "epoch": n} gives model.blocks.0.…. Nothing is stripped — what the file says is what you get.
A Pt is a handle: every copy is the same open checkpoint, and the last owner unmaps it. Nothing here is released by hand.
: i PT_MAX_DIMS 8Up to 8 dimensions is more than any real checkpoint tensor uses; the limit keeps the per-tensor record a fixed size.
: PtTensor: PtTensor {
String name
i dtype // PKS_* storage class
i ndim
i nelems
i data_off // absolute FILE offset of element 0
i nbytes // logical size: nelems x element size
i contiguous // 1 = row-major, 0 = a view we can only describe
i shape_at // where its 16 slots start in the Pt's `dims`: dims 0..7, then strides 8..15
}: PtImpl: PtImpl {
* u map
i map_size
b from_mmap
( Vec u ) buf // the bytes `map` points into, when they are not a mapping
( Vec PtTensor ) tensors
( Vec i ) dims // 16 slots per tensor (PtTensor.shape_at): one array, not one per tensor
}% Drop PtImpl% Drop PtImpl {
unsafe @ drop PtImpl p → v {
? . p from_mmap { : i32 _u ( munmap . p map . p map_size ) } {}
}
}The mapping is the raw resource: its last owner unmaps it, as pt_close did. The tensor table and the buffer go with the drop glue.
: Pt { s ctl }A Pt is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same open checkpoint, and the last owner releases it.
@ Pt_drop sink Pt h → v@ pt_none → PtA Pt that holds no checkpoint — for a slot that may be empty (a model whose checkpoint is some other format). pt_is_open tells them apart.
@ pt_is_open Pt p → b@ _pt_geti ( Vec i ) v i k → i@ pt_find Pt p__h s name → i@ pt_shape_str Pt p__h i idx → StringShape rendered as "2x3x4" — for CLI listings and error messages.
@ pt_close sink Pt p → vLet go of p now rather than at the end of its owner's scope; the last owner unmaps the checkpoint.
@ pt_dequant Pt p__h i idx → !( Vec u ) Stringpackages/torchpt/src/main.nu — the torchpt CLI.
torchpt info <file.pt> summary: tensors, dtypes, bytes torchpt tensors <file.pt> [prefix] one line per tensor torchpt stats <file.pt> <name> min / max / mean of one tensor torchpt head <file.pt> <name> [n] first n values torchpt dump <file.pt> every tensor, every value — the oracle-diff format
@ usage → i@ starts_with s hay s pre → b@ cmd_info Pt p → i@ cmd_tensors Pt p s prefix → i@ main → ipackages/torchpt/src/pickle.nu — a pickle protocol 0–5 reader, restricted to data.
Unpickling is famously arbitrary code execution: the format is a stack machine whose REDUCE opcode means "call this callable". This reader never calls anything. GLOBAL pushes the name of a callable, and REDUCE consults a fixed table of shapes it knows how to build as data — torch._utils._rebuild_tensor_v2, collections.OrderedDict and a couple of siblings. Everything else becomes an opaque PK_OTHER node. A hostile pickle can therefore make this reader produce a confusing tree, and nothing else.
The result is a value tree in flat arrays (no per-node allocation):
( pk_parse u p i n ) → !Pk String ( pk_free k ) → v early release (optional) ( pk_root k ) → i root node id ( __pk_kind k id ) → i PK_ ( pk_int k id ) → i PK_INT / PK_BOOL value ( pk_float k id ) → f PK_FLOAT value ( pk_str k id ) → s PK_STR text (BORROWED) ( __pk_len k id ) → i items (seq) or pairs (dict) ( pk_item k id j ) → i j-th item of a tuple/list ( pk_key k id j ) → i j-th key of a dict ( pk_val k id j ) → i j-th value of a dict ( pk_dict_get k id s key ) → i value node, -1 when absent
Tensors reduce to PK_TENSOR nodes; their fields come out through ( pk_tensor_storage k id ) → i storage id ( pk_tensor_offset k id ) → i element offset into the storage ( pk_tensor_ndim k id ) → i ( pk_tensor_dim k id j ) → i ( pk_tensor_stride k id j ) → i and a storage id resolves through ( pk_storage_dtype k sid ) → i PKS_* storage class ( pk_storage_key k sid ) → s the data/<key> filename (BORROWED) ( pk_storage_numel k sid ) → i
Everything is bounded before it is believed: a length read from the stream is checked against the bytes that remain, memo ids and stack depths are checked before use, and the opcode loop has a step budget proportional to the input, so a crafted pickle cannot spin, over-read or allocate on a number it merely claims.
A Pk is a handle: every copy is the same tree, and the last owner releases it. Nothing here is released by hand.
: i PK_NONE 0── node kinds ──────────────────────────────────────────────────────
: i PK_BOOL 1: i PK_INT 2: i PK_FLOAT 3: i PK_STR 4: i PK_BYTES 5: i PK_TUPLE 6: i PK_LIST 7: i PK_DICT 8: i PK_TENSOR 9: i PK_GLOBAL 10: i PK_PERSID 11: i PK_OTHER 12: i PK_MARK 13 // internal: stack sentinel, never reachable from the root@ pk_kind_name i kd → s: i PKS_UNKNOWN 0── torch storage classes ───────────────────────────────────────────
: i PKS_F64 1: i PKS_F32 2: i PKS_F16 3: i PKS_BF16 4: i PKS_I64 5: i PKS_I32 6: i PKS_I16 7: i PKS_I8 8: i PKS_U8 9: i PKS_BOOL 10@ pk_storage_class s name → i@ pk_storage_name i c → s@ pk_storage_esize i c → iBytes per element of a storage class; 0 when unknown.
: PkImpl: PkImpl {
( Vec i ) kind // per node
( Vec i ) va // per node: int / str-idx / kids-start / tensor-idx / storage-idx
( Vec i ) vb // per node: item or pair count, or byte length
( Vec i ) kids // child node ids, addressed by (va, vb)
( Vec String ) strs
( Vec u ) blob // BYTES payloads, addressed by (va, vb)
( Vec i ) stack
( Vec i ) memo
( Vec i ) stor // 3 per storage: class, key str-idx, numel
( Vec i ) tens // 4 per tensor: storage id, elem offset, ndim, dims-start
( Vec i ) dims // ndim dims then ndim strides, per tensor
( Vec i ) root // slot 0: root node id
( Vec String ) errs
}── the tree, in flat arrays ────────────────────────────────────────
Every field is a Vec, never a scalar: a NURL struct passed by value copies its scalar fields, so a counter living in the struct would silently not mutate. Vec handles are shared, so pushes are seen by the caller — which is why even the root id and the error live in one-element Vecs.
: Pk { s ctl }A Pk is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same tree, and the last owner releases it.
@ Pk_drop sink Pk h → v@ pk_free sink Pk k → vLet go of k now rather than at the end of its owner's scope.