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

#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"

Include dependency graph for n_pth.c:

Go to the source code of this file.

Data Structures
struct	PK_OBJ
	One pickle value. More...
struct	PK_VM
	The interpreter. More...

Macros
#define	LEFT (len - at)
#define	PTH_OBJECTS_MAX 4000000u
	Most values the interpreter makes, and most memo entries: a guard, far above any real checkpoint's.
#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_STORAGE_MAX ((size_t)4 * 1024u * 1024u * 1024u)
	Largest single storage read.

Enumerations
enum	PK_TYPE { 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 }
	What a pickle value is. More...

Functions
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
static float	element_value (const unsigned char p, int element, const char dtype)
	convert one stored element to a float
static uint64_t	le (const unsigned char *p, int n)
	read a little endian unsigned value
const N_PTH_TENSOR *	n_pth_find (const N_PTH pth, const char name)
	find a tensor by its path
void	n_pth_free (N_PTH **pth)
	release a checkpoint's tensors
N_PTH *	n_pth_load (const char *path)
	read every tensor of a checkpoint
static int	pk_add (PK_OBJ o, PK_OBJ item)
	add an item to a tuple, list or dictionary
static PK_OBJ *	pk_from_mark (PK_VM *vm, int type)
	pop the values above the top mark into a new container
static long	pk_mark (const PK_VM *vm)
	where the top mark is on the stack
static int	pk_memo_put (PK_VM vm, size_t index, PK_OBJ o)
	store a value in the memo
static PK_OBJ *	pk_new (PK_VM *vm, int type)
	make a value, owned by the interpreter
static PK_OBJ *	pk_pop (PK_VM *vm)
	pop a value
static int	pk_push (PK_VM vm, PK_OBJ o)
	push a value
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
static PK_OBJ *	pk_run (PK_VM vm, const unsigned char p, size_t len)
	run a pickle
static PK_OBJ *	pk_string (PK_VM vm, int type, const unsigned char p, size_t len)
	make a string value from bytes
static int	storage_element (const char name, const char *dtype)
	a storage type's element size, from its global's name
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

Detailed Description

Reading the tensors of a PyTorch checkpoint without Python.

Author: Castagnier Mickael

Version: 1.0

Date: 29/09/2026

Definition in file n_pth.c.

Data Structure Documentation

◆ PK_OBJ

struct PK_OBJ

One pickle value.

Dictionaries keep keys and values alternately in items.

Definition at line 66 of file n_pth.c.

Collaboration diagram for PK_OBJ:

Data Fields
size_t	cap	room for
const char *	dtype	a storage's type name, for the log
int	element	a storage's element size in bytes
double	f	a float's value
int64_t	i	an integer's value, a storage's element count, a tensor's offset
struct PK_OBJ **	items	items of a tuple, list or dictionary; a tensor's storage, then its size and stride tuples
size_t	n	how many
char *	s	a string, a global's "module name", a storage's key; owned
int	type	a PK_TYPE

◆ PK_VM

struct PK_VM

The interpreter.

Definition at line 89 of file n_pth.c.

Collaboration diagram for PK_VM:

Data Fields
PK_OBJ **	all	every value made, freed together
size_t	all_cap	room for
size_t	all_n	how many
PK_OBJ **	memo	the memo
size_t	memo_n	its size
size_t	sp	its depth
PK_OBJ **	stack	the stack
size_t	stack_cap	room for

Macro Definition Documentation

◆ LEFT

#define LEFT (len - at)

Referenced by pk_run().

◆ 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.

Definition at line 44 of file n_pth.c.

Referenced by pk_add(), pk_memo_put(), pk_new(), and pk_push().

◆ 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.

Definition at line 39 of file n_pth.c.

Referenced by n_pth_load().

◆ PTH_STORAGE_MAX

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

Largest single storage read.

Definition at line 41 of file n_pth.c.

Referenced by tensor_read().

Enumeration Type Documentation

◆ PK_TYPE

enum PK_TYPE

What a pickle value is.

Enumerator
PK_NONE
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

Definition at line 47 of file n_pth.c.

Function Documentation

◆ collect()

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

static

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