You can do anything in Python

Python isn’t my favorite language, but you can do anything with it. You can do almost anything in other programming languages too. But doing it in a high-level language makes it easier to look under the abstractions and see how they’re built.

One project that takes this to the max is tinygrad. For those who don’t know, it’s a neural network engine made by the legendary George Hotz and his company tinycorp. I remember seeing the quote below in their docs, and I think it sums up how tinygrad compares with other ML libraries:

If XLA is CISC, tinygrad is RISC

How are they achieving this? With pure Python with almost no dependencies, less than 30k LoC by the time of this post. Much of tinygrad’s compiler and control code is written in Python, not just its tensor operations. They even have an AMD GPU driver written in Python, writing their own compiler backends and soon it will remove the LLVM dependency. And yep, it’s getting very good.

The spark for this post came after I wrote a piece of my own GPT-2 implementation. Sebastian’s book has been teaching me a lot about the internals of the architecture and the computations happening inside LLMs. There’s a chapter that shows how to load the weights from the original TensorFlow checkpoint released with GPT-2, but I wanted to download and load them from Hugging Face. Since I know a bit of the tinygrad code, which has almost zero dependencies, I thought it might have a function to load weights from different formats without needing the safetensors library. And bingo, they had it. From there, I wrote my own based on it, also looking at the safetensors format spec.

safetensors is a format created by Hugging Face to avoid the arbitrary-code execution risk that comes with loading Pickle files. Pickle lets a class define a __reduce__() method that can specify a callable to run during unpickling. A malicious pickle file can use this to execute code when someone loads it. You can check an example in this Snyk blog post.

Safetensors stores tensor data and metadata, not executable Python objects, so loading a file doesn’t invoke code embedded in it. You still need to parse its header. The file has an 8-byte length prefix, a JSON header, and a packed tensor-data section. In outline:

0                    8                    8 + N
+--------------------+--------------------+-------------------------+
| header_len(N)      | JSON header        | tensor_data             |
+--------------------+--------------------+-------------------------+
                                          ^
                                          data_start

The first 8 bytes encode the header length, N, as an unsigned 64-bit little-endian integer. The next N bytes, at offsets [8, 8 + N), contain the JSON header with the tensor offsets, shapes, and dtypes. For example:

{
  "weight_1": {
    "dtype": "F32",
    "shape": [2],
    "data_offsets": [0, 8]
  }
}

Starting at byte 8 + N, the tensor data is packed contiguously. Each tensor’s data offsets are relative to the start of this data section. It’s easy to parse. This is a small example of reading it and loading the data into PyTorch tensors:

import torch
import json
from pathlib import Path

def load_safetensors(path: Path) -> dict[str, torch.Tensor]:
    SAFETENSORS_DTYPES = {
        "F64": torch.float64,
        "F32": torch.float32,
        "F16": torch.float16,
        "BF16": torch.bfloat16,
    }

    # with a context manager, read the file
    with path.open("rb") as file:
        # read the first 8 bytes to get the header length. 
        # file.read() moves the pointer to the number of bytes provided
        length_bytes = file.read(8)
        if len(length_bytes) != 8:
            raise ValueError("Truncated safetensors header")
        # parse length_bytes as a 8 byte (u64) integer represented in little-endian
        header_size = int.from_bytes(length_bytes, byteorder="little")
        # read header_size bytes and parse as a JSON, containing the weights metadata
        header = json.loads(file.read(header_size))
        # return the current memory position
        data_start = file.tell()
        tensors = {}

        for name, metadata in header.items():
            # free-form string-to-string map, can be ignored
            if name == "__metadata__":
                continue

            # read the weight from the offsets provided
            start, end = metadata["data_offsets"]
            file.seek(data_start + start)
            # convert into a PyTorch tensor with the configured dtype and shape
            raw_tensor = bytearray(file.read(end - start))
            dtype = SAFETENSORS_DTYPES[metadata["dtype"]]
            tensors[name] = torch.frombuffer(raw_tensor, dtype=dtype).reshape(
                metadata["shape"]
            )

    return tensors

A lot of programmers would rather just think that parsing a data format would be too hard and should not concern with it, but from writing this code you avoid suppy chain attacks that happens when with multiple open source libraries. And thinking this is hard is basically saying you think CS 101 is hard. Computers read bytes and operate under those, there’s not much more than this.

Of course you can do this in C and make it way faster, but doing it in Python made it easy to follow every step. So yeah, whenever someone says Python is just for high-level wrappers, you can show them tinygrad and convince them they’re wrong.