src/n_pth.c

Reading the tensors of a PyTorch checkpoint without Python. More...

Structures

	Name
struct	PK_VM The interpreter.
struct	PK_OBJ One pickle value.

Types

	Name
enum	PK_TYPE { PK_TUPLE, PK_TENSOR, PK_STR, PK_STORAGE, PK_OPAQUE, PK_NONE = 0, PK_MARK, PK_LIST, PK_INT, PK_GLOBAL, PK_FLOAT, PK_DICT, PK_BYTES, PK_BOOL} What a pickle value is.

Functions

	Name
int	tensor_read(const char * path, const char * dir, const PK_OBJ * t, N_PTH_TENSOR * out) read a tensor's values out of its storage file
int	storage_element(const char * name, const char ** dtype) a storage type's element size, from its global's name
PK_OBJ *	pk_string(PK_VM * vm, int type, const unsigned char * p, size_t len) make a string value from bytes
PK_OBJ *	pk_run(PK_VM * vm, const unsigned char * p, size_t len) run a pickle
PK_OBJ *	pk_reduce(PK_VM * vm, const PK_OBJ * callable, PK_OBJ * args) the call a REDUCE or NEWOBJ makes, done as data
int	pk_push(PK_VM * vm, PK_OBJ * o) push a value
PK_OBJ *	pk_pop(PK_VM * vm) pop a value
PK_OBJ *	pk_new(PK_VM * vm, int type) make a value, owned by the interpreter
int	pk_memo_put(PK_VM * vm, size_t index, PK_OBJ * o) store a value in the memo
long	pk_mark(const PK_VM * vm) where the top mark is on the stack
PK_OBJ *	pk_from_mark(PK_VM * vm, int type) pop the values above the top mark into a new container
int	pk_add(PK_OBJ * o, PK_OBJ * item) add an item to a tuple, list or dictionary
uint64_t	le(const unsigned char * p, int n) read a little endian unsigned value
float	element_value(const unsigned char * p, int element, const char * dtype) convert one stored element to a float
int	collect(const char * path, const char * dir, const PK_OBJ * o, const char * prefix, N_PTH * pth, int depth) collect the tensors under a dictionary, depth first

Defines

	Name
	PTH_STORAGE_MAX Largest single storage read.
	PTH_PICKLE_MAX Largest data.pkl read: a state dictionary's is a few kilobytes, a large model's a few hundred.
	PTH_OBJECTS_MAX Most values the interpreter makes, and most memo entries: a guard, far above any real checkpoint's.
	LEFT

Detailed Description

Reading the tensors of a PyTorch checkpoint without Python.

Author: Castagnier Mickael

Version: 1.0

Date: 29/09/2026

Types Documentation

enum PK_TYPE

Enumerator	Value	Description
PK_TUPLE
PK_TENSOR
PK_STR
PK_STORAGE
PK_OPAQUE
PK_NONE	0
PK_MARK
PK_LIST
PK_INT
PK_GLOBAL
PK_FLOAT
PK_DICT
PK_BYTES
PK_BOOL

What a pickle value is.

Functions Documentation

function tensor_read

static int tensor_read(
    const char * path,
    const char * dir,
    const PK_OBJ * t,
    N_PTH_TENSOR * out
)

read a tensor's values out of its storage file

Parameters:

• path the archive
• dir the folder data.pkl is in, "archive/"
• t the tensor value
• out receives the tensor

Return: 0, or -1

function storage_element

static int storage_element(
    const char * name,
    const char ** dtype
)

a storage type's element size, from its global's name

Parameters:

• name the global, "torch FloatStorage"
• dtype receives a short name

Return: 4, 2 or 8, or 0 for a type not read

function pk_string

static PK_OBJ * pk_string(
    PK_VM * vm,
    int type,
    const unsigned char * p,
    size_t len
)

make a string value from bytes

Parameters:

• vm the interpreter
• type PK_STR or PK_BYTES
• p the bytes
• len how many

Return: the value, or NULL

function pk_run

static PK_OBJ * pk_run(
    PK_VM * vm,
    const unsigned char * p,
    size_t len
)

run a pickle

Parameters:

• vm the interpreter
• p the pickle
• len its length

Return: the value it leaves, or NULL

function pk_reduce

static PK_OBJ * pk_reduce(
    PK_VM * vm,
    const PK_OBJ * callable,
    PK_OBJ * args
)

the call a REDUCE or NEWOBJ makes, done as data

Parameters:

• vm the interpreter
• callable what is called
• args its arguments, a tuple

Return: the result, or NULL

function pk_push

static int pk_push(
    PK_VM * vm,
    PK_OBJ * o
)

push a value

Parameters:

• vm the interpreter
• o the value; NULL fails

Return: 0, or -1

function pk_pop

static PK_OBJ * pk_pop(
    PK_VM * vm
)

pop a value

Parameters:

• vm the interpreter

Return: the value, or NULL on an empty stack or at a mark

function pk_new

static PK_OBJ * pk_new(
    PK_VM * vm,
    int type
)

make a value, owned by the interpreter

Parameters:

• vm the interpreter
• type its PK_TYPE

Return: the value, or NULL

function pk_memo_put

static int pk_memo_put(
    PK_VM * vm,
    size_t index,
    PK_OBJ * o
)

store a value in the memo

Parameters:

• vm the interpreter
• index where
• o the value

Return: 0, or -1

function pk_mark

static long pk_mark(
    const PK_VM * vm
)

where the top mark is on the stack

Parameters:

• vm the interpreter

Return: its index, or -1 when there is none

function pk_from_mark

static PK_OBJ * pk_from_mark(
    PK_VM * vm,
    int type
)

pop the values above the top mark into a new container

Parameters:

• vm the interpreter
• type PK_TUPLE or PK_LIST

Return: the container, or NULL

function pk_add

static int pk_add(
    PK_OBJ * o,
    PK_OBJ * item
)

add an item to a tuple, list or dictionary

Parameters:

• o the container
• item the item

Return: 0, or -1

function le

static uint64_t le(
    const unsigned char * p,
    int n
)

read a little endian unsigned value

Parameters:

• p the bytes
• n how many, up to 8

Return: the value

function element_value

static float element_value(
    const unsigned char * p,
    int element,
    const char * dtype
)

convert one stored element to a float

Parameters:

• p the element's bytes
• element its size
• dtype its type name

Return: the value

function collect

static int collect(
    const char * path,
    const char * dir,
    const PK_OBJ * o,
    const char * prefix,
    N_PTH * pth,
    int depth
)

collect the tensors under a dictionary, depth first

Parameters:

• path the archive
• dir the folder data.pkl is in
• o a value
• prefix its path so far
• pth receives the tensors
• depth how deep, a guard against a dictionary inside itself

Return: 0, or -1

Macros Documentation

define PTH_STORAGE_MAX

#define PTH_STORAGE_MAX ((size_t)4 * 1024u * 1024u * 1024u)

Largest single storage read.

define PTH_PICKLE_MAX

#define PTH_PICKLE_MAX (64u * 1024u * 1024u)

Largest data.pkl read: a state dictionary's is a few kilobytes, a large model's a few hundred.

define PTH_OBJECTS_MAX

#define PTH_OBJECTS_MAX 4000000u

Most values the interpreter makes, and most memo entries: a guard, far above any real checkpoint's.

define LEFT

#define LEFT (len - at)

Source code

/*
 * Nilorea Library
 * Copyright (C) 2005-2026 Castagnier Mickael
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 */


#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "nilorea/n_log.h"
#include "nilorea/n_pth.h"
#include "nilorea/n_zip.h"

#define PTH_PICKLE_MAX (64u * 1024u * 1024u)
#define PTH_STORAGE_MAX ((size_t)4 * 1024u * 1024u * 1024u)
#define PTH_OBJECTS_MAX 4000000u

typedef enum {
    PK_NONE = 0,
    PK_BOOL,
    PK_INT,
    PK_FLOAT,
    PK_STR,
    PK_BYTES,
    PK_TUPLE,
    PK_LIST,
    PK_DICT,
    PK_GLOBAL,
    PK_STORAGE,
    PK_TENSOR,
    PK_OPAQUE,
    PK_MARK
} PK_TYPE;

typedef struct PK_OBJ {
    int type;
    int64_t i;
    double f;
    char* s;
    int element;
    const char* dtype;
    struct PK_OBJ** items;
    size_t n;
    size_t cap;
} PK_OBJ;

typedef struct {
    PK_OBJ** all;
    size_t all_n;
    size_t all_cap;
    PK_OBJ** stack;
    size_t sp;
    size_t stack_cap;
    PK_OBJ** memo;
    size_t memo_n;
} PK_VM;

static PK_OBJ* pk_new(PK_VM* vm, int type) {
    PK_OBJ* o = NULL;

    if (vm->all_n >= PTH_OBJECTS_MAX) return NULL;
    if (vm->all_n == vm->all_cap) {
        size_t cap = vm->all_cap ? vm->all_cap * 2 : 256;
        PK_OBJ** grown = (PK_OBJ**)realloc((void*)vm->all, cap * sizeof(*grown));

        if (!grown) return NULL;
        vm->all = grown;
        vm->all_cap = cap;
    }
    o = (PK_OBJ*)calloc(1, sizeof(*o));
    if (!o) return NULL;
    o->type = type;
    vm->all[vm->all_n++] = o;
    return o;
}

static int pk_add(PK_OBJ* o, PK_OBJ* item) {
    if (o->n == o->cap) {
        size_t cap = o->cap ? o->cap * 2 : 4;
        PK_OBJ** grown = NULL;

        if (cap > PTH_OBJECTS_MAX) return -1;
        grown = (PK_OBJ**)realloc((void*)o->items, cap * sizeof(*grown));
        if (!grown) return -1;
        o->items = grown;
        o->cap = cap;
    }
    o->items[o->n++] = item;
    return 0;
}

static int pk_push(PK_VM* vm, PK_OBJ* o) {
    if (!o) return -1;
    if (vm->sp == vm->stack_cap) {
        size_t cap = vm->stack_cap ? vm->stack_cap * 2 : 64;
        PK_OBJ** grown = NULL;

        if (cap > PTH_OBJECTS_MAX) return -1;
        grown = (PK_OBJ**)realloc((void*)vm->stack, cap * sizeof(*grown));
        if (!grown) return -1;
        vm->stack = grown;
        vm->stack_cap = cap;
    }
    vm->stack[vm->sp++] = o;
    return 0;
}

static PK_OBJ* pk_pop(PK_VM* vm) {
    if (!vm->sp || vm->stack[vm->sp - 1]->type == PK_MARK) return NULL;
    return vm->stack[--vm->sp];
}

static long pk_mark(const PK_VM* vm) {
    size_t i = vm->sp;

    while (i > 0) {
        if (vm->stack[i - 1]->type == PK_MARK) return (long)(i - 1);
        i--;
    }
    return -1;
}

static PK_OBJ* pk_string(PK_VM* vm, int type, const unsigned char* p, size_t len) {
    PK_OBJ* o = pk_new(vm, type);

    if (!o) return NULL;
    o->s = (char*)malloc(len + 1);
    if (!o->s) return NULL;
    memcpy(o->s, p, len);
    o->s[len] = '\0';
    o->i = (int64_t)len;
    return o;
}

static int pk_memo_put(PK_VM* vm, size_t index, PK_OBJ* o) {
    if (!o || index >= PTH_OBJECTS_MAX) return -1;
    if (index >= vm->memo_n) {
        size_t n = index + 1 > vm->memo_n * 2 ? index + 1 : vm->memo_n * 2;
        PK_OBJ** grown = (PK_OBJ**)realloc((void*)vm->memo, n * sizeof(*grown));

        if (!grown) return -1;
        memset((void*)(grown + vm->memo_n), 0, (n - vm->memo_n) * sizeof(*grown));
        vm->memo = grown;
        vm->memo_n = n;
    }
    vm->memo[index] = o;
    return 0;
}

static uint64_t le(const unsigned char* p, int n) {
    uint64_t v = 0;
    int k = 0;

    for (k = n - 1; k >= 0; k--) v = (v << 8) | p[k];
    return v;
}

static int storage_element(const char* name, const char** dtype) {
    static const struct {
        const char* global;
        const char* dtype;
        int size;
    } known[] = {{"torch FloatStorage", "float32", 4},
                 {"torch HalfStorage", "float16", 2},
                 {"torch BFloat16Storage", "bfloat16", 2},
                 {"torch DoubleStorage", "float64", 8},
                 {NULL, NULL, 0}};
    int i = 0;

    for (i = 0; known[i].global; i++) {
        if (strcmp(name, known[i].global) == 0) {
            *dtype = known[i].dtype;
            return known[i].size;
        }
    }
    *dtype = name;
    return 0;
}

static PK_OBJ* pk_reduce(PK_VM* vm, const PK_OBJ* callable, PK_OBJ* args) {
    if (callable->type == PK_GLOBAL && strcmp(callable->s, "collections OrderedDict") == 0) return pk_new(vm, PK_DICT);
    if (callable->type == PK_GLOBAL && (strcmp(callable->s, "torch._utils _rebuild_tensor_v2") == 0 ||
                                        strcmp(callable->s, "torch._utils _rebuild_tensor") == 0)) {
        PK_OBJ* t = NULL;

        /* (storage, offset, size, stride, ...) */
        if (args->type != PK_TUPLE || args->n < 4 || args->items[0]->type != PK_STORAGE ||
            args->items[1]->type != PK_INT || args->items[2]->type != PK_TUPLE || args->items[3]->type != PK_TUPLE)
            return pk_new(vm, PK_OPAQUE);
        t = pk_new(vm, PK_TENSOR);
        if (!t) return NULL;
        t->i = args->items[1]->i;
        if (pk_add(t, args->items[0]) || pk_add(t, args->items[2]) || pk_add(t, args->items[3])) return NULL;
        return t;
    }
    if (callable->type == PK_GLOBAL && strcmp(callable->s, "torch._utils _rebuild_parameter") == 0 &&
        args->type == PK_TUPLE && args->n >= 1 && args->items[0]->type == PK_TENSOR)
        return args->items[0];
    return pk_new(vm, PK_OPAQUE);
}

static PK_OBJ* pk_from_mark(PK_VM* vm, int type) {
    long m = pk_mark(vm);
    PK_OBJ* o = NULL;
    size_t k = 0;

    if (m < 0) return NULL;
    o = pk_new(vm, type);
    if (!o) return NULL;
    for (k = (size_t)m + 1; k < vm->sp; k++)
        if (pk_add(o, vm->stack[k])) return NULL;
    vm->sp = (size_t)m;
    return o;
}

static PK_OBJ* pk_run(PK_VM* vm, const unsigned char* p, size_t len) {
    size_t at = 0;

/* How many bytes are left, for the checks. */
#define LEFT (len - at)
    while (at < len) {
        unsigned char op = p[at++];
        PK_OBJ *a = NULL, *b = NULL, *c = NULL;

        switch (op) {
            case 0x80: /* PROTO */
                if (LEFT < 1) return NULL;
                at += 1;
                break;
            case 0x95: /* FRAME */
                if (LEFT < 8) return NULL;
                at += 8;
                break;
            case '.': /* STOP */
                return pk_pop(vm);
            case '(': /* MARK */
                if (pk_push(vm, pk_new(vm, PK_MARK))) return NULL;
                break;
            case '0': /* POP */
                if (!vm->sp) return NULL;
                vm->sp--;
                break;
            case '1': { /* POP_MARK */
                long m = pk_mark(vm);

                if (m < 0) return NULL;
                vm->sp = (size_t)m;
                break;
            }
            case '}': /* EMPTY_DICT */
                if (pk_push(vm, pk_new(vm, PK_DICT))) return NULL;
                break;
            case ']': /* EMPTY_LIST */
                if (pk_push(vm, pk_new(vm, PK_LIST))) return NULL;
                break;
            case ')': /* EMPTY_TUPLE */
                if (pk_push(vm, pk_new(vm, PK_TUPLE))) return NULL;
                break;
            case 't': /* TUPLE */
                if (pk_push(vm, pk_from_mark(vm, PK_TUPLE))) return NULL;
                break;
            case 'l': /* LIST */
                if (pk_push(vm, pk_from_mark(vm, PK_LIST))) return NULL;
                break;
            case 0x85:   /* TUPLE1 */
            case 0x86:   /* TUPLE2 */
            case 0x87: { /* TUPLE3 */
                size_t n = (size_t)(op - 0x84), k = 0;

                if (vm->sp < n) return NULL;
                a = pk_new(vm, PK_TUPLE);
                if (!a) return NULL;
                for (k = vm->sp - n; k < vm->sp; k++) {
                    if (vm->stack[k]->type == PK_MARK || pk_add(a, vm->stack[k])) return NULL;
                }
                vm->sp -= n;
                if (pk_push(vm, a)) return NULL;
                break;
            }
            case 'a': /* APPEND */
                b = pk_pop(vm);
                if (!b || !vm->sp || vm->stack[vm->sp - 1]->type != PK_LIST || pk_add(vm->stack[vm->sp - 1], b))
                    return NULL;
                break;
            case 'e': { /* APPENDS */
                long m = pk_mark(vm);
                size_t k = 0;

                if (m < 1 || vm->stack[m - 1]->type != PK_LIST) return NULL;
                for (k = (size_t)m + 1; k < vm->sp; k++)
                    if (pk_add(vm->stack[m - 1], vm->stack[k])) return NULL;
                vm->sp = (size_t)m;
                break;
            }
            case 'd': { /* DICT */
                long m = pk_mark(vm);
                size_t k = 0;

                if (m < 0 || (vm->sp - (size_t)m - 1) % 2) return NULL;
                a = pk_new(vm, PK_DICT);
                if (!a) return NULL;
                for (k = (size_t)m + 1; k < vm->sp; k++)
                    if (pk_add(a, vm->stack[k])) return NULL;
                vm->sp = (size_t)m;
                if (pk_push(vm, a)) return NULL;
                break;
            }
            case 's': /* SETITEM */
                b = pk_pop(vm);
                a = pk_pop(vm);
                if (!a || !b || !vm->sp || vm->stack[vm->sp - 1]->type != PK_DICT) return NULL;
                if (pk_add(vm->stack[vm->sp - 1], a) || pk_add(vm->stack[vm->sp - 1], b)) return NULL;
                break;
            case 'u': { /* SETITEMS */
                long m = pk_mark(vm);

                if (m < 1 || (vm->sp - (size_t)m - 1) % 2) return NULL;
                /* Items set on an opaque value, an object of a class this
                 * does not build, are dropped with it. */
                if (vm->stack[m - 1]->type == PK_DICT) {
                    for (size_t k = (size_t)m + 1; k < vm->sp; k++)
                        if (pk_add(vm->stack[m - 1], vm->stack[k])) return NULL;
                } else if (vm->stack[m - 1]->type != PK_OPAQUE) {
                    return NULL;
                }
                vm->sp = (size_t)m;
                break;
            }
            case 'J': /* BININT */
                if (LEFT < 4 || !(a = pk_new(vm, PK_INT))) return NULL;
                a->i = (int32_t)(uint32_t)le(p + at, 4);
                at += 4;
                if (pk_push(vm, a)) return NULL;
                break;
            case 'K': /* BININT1 */
                if (LEFT < 1 || !(a = pk_new(vm, PK_INT))) return NULL;
                a->i = p[at++];
                if (pk_push(vm, a)) return NULL;
                break;
            case 'M': /* BININT2 */
                if (LEFT < 2 || !(a = pk_new(vm, PK_INT))) return NULL;
                a->i = (int64_t)le(p + at, 2);
                at += 2;
                if (pk_push(vm, a)) return NULL;
                break;
            case 0x8a: { /* LONG1 */
                size_t n = 0;

                if (LEFT < 1) return NULL;
                n = p[at++];
                if (n > 8 || LEFT < n || !(a = pk_new(vm, PK_INT))) return NULL;
                a->i = n ? (int64_t)le(p + at, (int)n) : 0;
                /* Sign extend what was shorter than eight bytes. */
                if (n && n < 8 && (p[at + n - 1] & 0x80)) a->i -= (int64_t)1 << (8 * n);
                at += n;
                if (pk_push(vm, a)) return NULL;
                break;
            }
            case 'G': { /* BINFLOAT, big endian */
                uint64_t bits = 0;
                int k = 0;

                if (LEFT < 8 || !(a = pk_new(vm, PK_FLOAT))) return NULL;
                for (k = 0; k < 8; k++) bits = (bits << 8) | p[at + (size_t)k];
                memcpy(&a->f, &bits, sizeof(a->f));
                at += 8;
                if (pk_push(vm, a)) return NULL;
                break;
            }
            case 'N': /* NONE */
                if (pk_push(vm, pk_new(vm, PK_NONE))) return NULL;
                break;
            case 0x88: /* NEWTRUE */
            case 0x89: /* NEWFALSE */
                if (!(a = pk_new(vm, PK_BOOL))) return NULL;
                a->i = op == 0x88;
                if (pk_push(vm, a)) return NULL;
                break;
            case 'X':    /* BINUNICODE */
            case 'T':    /* BINSTRING */
            case 'B':    /* BINBYTES */
            case 0x8c:   /* SHORT_BINUNICODE */
            case 'U':    /* SHORT_BINSTRING */
            case 'C':    /* SHORT_BINBYTES */
            case 0x8d:   /* BINUNICODE8 */
            case 0x8e: { /* BINBYTES8 */
                int width = (op == 0x8c || op == 'U' || op == 'C') ? 1 : (op == 0x8d || op == 0x8e) ? 8
                                                                                                    : 4;
                uint64_t n = 0;

                if (LEFT < (size_t)width) return NULL;
                n = le(p + at, width);
                at += (size_t)width;
                if (n > LEFT) return NULL;
                if (pk_push(vm, pk_string(vm, (op == 'B' || op == 'C' || op == 0x8e) ? PK_BYTES : PK_STR, p + at,
                                          (size_t)n)))
                    return NULL;
                at += (size_t)n;
                break;
            }
            case 'c': { /* GLOBAL: "module\nname\n" */
                const unsigned char* first = memchr(p + at, '\n', LEFT);
                const unsigned char* second = first ? memchr(first + 1, '\n', (size_t)(p + len - first - 1)) : NULL;

                if (!first || !second) return NULL;
                a = pk_string(vm, PK_GLOBAL, p + at, (size_t)(second - (p + at)));
                if (!a) return NULL;
                a->s[first - (p + at)] = ' ';
                at = (size_t)(second - p) + 1;
                if (pk_push(vm, a)) return NULL;
                break;
            }
            case 0x93: { /* STACK_GLOBAL */
                size_t n = 0;

                b = pk_pop(vm);
                a = pk_pop(vm);
                if (!a || !b || a->type != PK_STR || b->type != PK_STR) return NULL;
                n = strlen(a->s) + strlen(b->s) + 2;
                c = pk_new(vm, PK_GLOBAL);
                if (!c || !(c->s = (char*)malloc(n))) return NULL;
                snprintf(c->s, n, "%s %s", a->s, b->s);
                if (pk_push(vm, c)) return NULL;
                break;
            }
            case 'q': /* BINPUT */
                if (LEFT < 1 || !vm->sp || pk_memo_put(vm, p[at], vm->stack[vm->sp - 1])) return NULL;
                at += 1;
                break;
            case 'r': /* LONG_BINPUT */
                if (LEFT < 4 || !vm->sp || pk_memo_put(vm, (size_t)le(p + at, 4), vm->stack[vm->sp - 1]))
                    return NULL;
                at += 4;
                break;
            case 0x94: /* MEMOIZE */
                if (!vm->sp) return NULL;
                {
                    size_t k = 0;

                    /* The next free index: one past the highest in use. */
                    for (k = vm->memo_n; k > 0 && !vm->memo[k - 1]; k--) {
                    }
                    if (pk_memo_put(vm, k, vm->stack[vm->sp - 1])) return NULL;
                }
                break;
            case 'h':   /* BINGET */
            case 'j': { /* LONG_BINGET */
                size_t index = 0;

                if (LEFT < (op == 'h' ? 1u : 4u)) return NULL;
                index = op == 'h' ? p[at] : (size_t)le(p + at, 4);
                at += op == 'h' ? 1 : 4;
                if (index >= vm->memo_n || !vm->memo[index] || pk_push(vm, vm->memo[index])) return NULL;
                break;
            }
            case 'Q': { /* BINPERSID: ("storage", type, key, location, numel) */
                const char* dtype = NULL;

                a = pk_pop(vm);
                if (!a || a->type != PK_TUPLE || a->n < 5 || a->items[1]->type != PK_GLOBAL ||
                    a->items[2]->type != PK_STR || a->items[4]->type != PK_INT)
                    return NULL;
                b = pk_new(vm, PK_STORAGE);
                if (!b) return NULL;
                b->element = storage_element(a->items[1]->s, &dtype);
                b->dtype = dtype;
                b->i = a->items[4]->i;
                b->s = (char*)malloc(strlen(a->items[2]->s) + 1);
                if (!b->s) return NULL;
                memcpy(b->s, a->items[2]->s, strlen(a->items[2]->s) + 1);
                if (pk_push(vm, b)) return NULL;
                break;
            }
            case 'R':  /* REDUCE */
            case 0x81: /* NEWOBJ */
                b = pk_pop(vm);
                a = pk_pop(vm);
                if (!a || !b || b->type != PK_TUPLE) return NULL;
                if (pk_push(vm, pk_reduce(vm, a, b))) return NULL;
                break;
            case 'b': /* BUILD: state set on an object, not kept */
                if (!pk_pop(vm) || !vm->sp) return NULL;
                break;
            default:
                n_log(LOG_ERR, "pth: pickle opcode 0x%02x at %zu is not read", op, at - 1);
                return NULL;
        }
    }
#undef LEFT
    return NULL;
}

static float element_value(const unsigned char* p, int element, const char* dtype) {
    uint32_t bits = 0;
    float f = 0.0f;

    if (element == 4) {
        bits = (uint32_t)le(p, 4);
    } else if (element == 8) {
        uint64_t b64 = le(p, 8);
        double d = 0.0;

        memcpy(&d, &b64, sizeof(d));
        return (float)d;
    } else if (strcmp(dtype, "bfloat16") == 0) {
        bits = (uint32_t)le(p, 2) << 16;
    } else {
        /* IEEE half to single. */
        uint32_t h = (uint32_t)le(p, 2), sign = (h & 0x8000u) << 16, exp = (h >> 10) & 0x1fu, man = h & 0x3ffu;

        if (exp == 0) {
            if (man == 0) {
                bits = sign;
            } else {
                exp = 127 - 15 + 1;
                while (!(man & 0x400u)) {
                    man <<= 1;
                    exp--;
                }
                bits = sign | (exp << 23) | ((man & 0x3ffu) << 13);
            }
        } else if (exp == 31) {
            bits = sign | 0x7f800000u | (man << 13);
        } else {
            bits = sign | ((exp - 15 + 127) << 23) | (man << 13);
        }
    }
    memcpy(&f, &bits, sizeof(f));
    return f;
}

static int tensor_read(const char* path, const char* dir, const PK_OBJ* t, N_PTH_TENSOR* out) {
    const PK_OBJ* storage = t->items[0];
    const PK_OBJ* size = t->items[1];
    const PK_OBJ* stride = t->items[2];
    unsigned char* raw = NULL;
    size_t raw_len = 0, i = 0, count = 1, span = 0;
    char member[512];
    int d = 0;

    if (!storage->element) {
        n_log(LOG_ERR, "pth: %s has storage of type %s, not read", out->name, storage->dtype);
        return -1;
    }
    if (size->n > N_PTH_MAX_DIMS || size->n != stride->n || t->i < 0) return -1;
    out->dims = (int)size->n;
    for (d = 0; d < out->dims; d++) {
        if (size->items[d]->type != PK_INT || stride->items[d]->type != PK_INT || size->items[d]->i < 0 ||
            stride->items[d]->i < 0)
            return -1;
        out->shape[d] = size->items[d]->i;
        if (out->shape[d] && count > (SIZE_MAX / sizeof(float)) / (size_t)out->shape[d]) return -1;
        count *= (size_t)out->shape[d];
        if (out->shape[d]) span += (size_t)(out->shape[d] - 1) * (size_t)stride->items[d]->i;
    }
    /* The last element read must be inside the storage. */
    if (count && (size_t)t->i + span >= (size_t)storage->i) return -1;
    if (snprintf(member, sizeof(member), "%sdata/%s", dir, storage->s) >= (int)sizeof(member)) return -1;
    if (n_zip_read(path, member, PTH_STORAGE_MAX, &raw, &raw_len) != 0) return -1;
    if (raw_len < (size_t)storage->i * (size_t)storage->element) {
        free(raw);
        return -1;
    }
    out->count = count;
    out->data = (float*)malloc((count ? count : 1) * sizeof(float));
    if (!out->data) {
        free(raw);
        return -1;
    }
    /* Row major, each index turned into the storage's position by the
     * strides: a view is gathered, a contiguous tensor read in order. */
    for (i = 0; i < count; i++) {
        size_t rest = i, at = (size_t)t->i;

        for (d = out->dims - 1; d >= 0; d--) {
            size_t n = (size_t)out->shape[d];

            at += (rest % n) * (size_t)stride->items[d]->i;
            rest /= n;
        }
        out->data[i] = element_value(raw + at * (size_t)storage->element, storage->element, storage->dtype);
    }
    free(raw);
    return 0;
}

static int collect(const char* path, const char* dir, const PK_OBJ* o, const char* prefix, N_PTH* pth, int depth) {
    size_t k = 0;

    if (depth > 32) return -1;
    if (o->type == PK_TENSOR) {
        N_PTH_TENSOR* grown = (N_PTH_TENSOR*)realloc(pth->tensors, (pth->count + 1) * sizeof(*grown));
        N_PTH_TENSOR* t = NULL;

        if (!grown) return -1;
        pth->tensors = grown;
        t = &pth->tensors[pth->count];
        memset(t, 0, sizeof(*t));
        t->name = (char*)malloc(strlen(prefix) + 1);
        if (!t->name) return -1;
        memcpy(t->name, prefix, strlen(prefix) + 1);
        pth->count++;
        return tensor_read(path, dir, o, t);
    }
    if (o->type != PK_DICT) return 0;
    for (k = 0; k + 1 < o->n; k += 2) {
        const PK_OBJ* key = o->items[k];
        char name[1024];

        if (key->type != PK_STR) continue;
        if (snprintf(name, sizeof(name), "%s%s%s", prefix, prefix[0] ? "." : "", key->s) >= (int)sizeof(name))
            return -1;
        if (collect(path, dir, o->items[k + 1], name, pth, depth + 1) != 0) return -1;
    }
    return 0;
}

N_PTH* n_pth_load(const char* path) {
    PK_VM vm;
    N_PTH* pth = NULL;
    unsigned char* pickle = NULL;
    size_t pickle_len = 0, k = 0;
    char member[512], dir[512];
    const PK_OBJ* top = NULL;
    const char* slash = NULL;

    if (!path) return NULL;
    if (n_zip_find_ending(path, "data.pkl", member, sizeof(member)) != 0) {
        n_log(LOG_ERR, "pth: %s is not a checkpoint this reads (no data.pkl)", path);
        return NULL;
    }
    /* The storages are under data/ beside data.pkl. */
    slash = strrchr(member, '/');
    snprintf(dir, sizeof(dir), "%.*s", slash ? (int)(slash - member + 1) : 0, member);
    if (n_zip_read(path, member, PTH_PICKLE_MAX, &pickle, &pickle_len) != 0) return NULL;
    memset(&vm, 0, sizeof(vm));
    top = pk_run(&vm, pickle, pickle_len);
    free(pickle);
    if (top) {
        pth = (N_PTH*)calloc(1, sizeof(*pth));
        if (pth && collect(path, dir, top, "", pth, 0) != 0) n_pth_free(&pth);
    } else {
        n_log(LOG_ERR, "pth: the pickle in %s could not be read", path);
    }
    for (k = 0; k < vm.all_n; k++) {
        free(vm.all[k]->s);
        free((void*)vm.all[k]->items);
        free(vm.all[k]);
    }
    free((void*)vm.all);
    free((void*)vm.stack);
    free((void*)vm.memo);
    return pth;
}

const N_PTH_TENSOR* n_pth_find(const N_PTH* pth, const char* name) {
    size_t k = 0;

    for (k = 0; pth && name && k < pth->count; k++)
        if (strcmp(pth->tensors[k].name, name) == 0) return &pth->tensors[k];
    return NULL;
}

void n_pth_free(N_PTH** pth) {
    size_t k = 0;

    if (!pth || !*pth) return;
    for (k = 0; k < (*pth)->count; k++) {
        free((*pth)->tensors[k].name);
        free((*pth)->tensors[k].data);
    }
    free((*pth)->tensors);
    free(*pth);
    *pth = NULL;
}