Polygrad Python

The Python API for Polygrad, a tensor library and JIT compiler built around a C11 port of tinygrad. It provides automatic differentiation, neural-network layers, optimizers and NumPy data exchange. Models saved in Python can load in JavaScript without their original Python classes.

Project overview and shared C/runtime reference | JavaScript guide

Contents

• Install
• Quickstart: Tensors
• Quickstart: Models
• Working With Tensors
• Training
• Models
• Devices And Runtimes
• JIT And Compile
• Custom Kernels
• Package Integration
• API Overview
• Troubleshooting

Install

pip install polygrad

Requirements:

• Linux
• Python 3.9 or newer (installed wheels tested on CPython 3.9-3.13)
• NumPy
• A C compiler on PATH for CPU execution: clang recommended, GCC fallback; generated float16 kernels require clang's __fp16 support

CPython 3.9-3.13 on x86_64 Linux with glibc 2.27+ can install a prebuilt wheel. Other Linux configurations build from source and require a C compiler and Python development headers. Wheels do not remove the CPU backend's runtime compiler requirement; use DEV=X86 (x86_64) or DEV=INTERP to run without one.

Optional model-loading dependency:

pip install huggingface_hub

For a source checkout, see building from source.

Quickstart: Tensors

from polygrad import Tensor

a = Tensor.rand(3, 4)
b = Tensor.rand(4, 5)
c = (a @ b).softmax(-1)

print(c.numpy())

Autograd:

from polygrad import Tensor

x = Tensor([1.0, 2.0, 3.0])

loss = (x * x).sum()
loss.backward()

print(x.grad.numpy())  # [2. 4. 6.]

Quickstart: Models

Fit y = 3x + 2 in Python, save its graph and weights, then load it in JavaScript without redefining the model.

from polygrad import Model, Tensor

class Linear:
    def __init__(self):
        self.a = Tensor([0.0])
        self.b = Tensor([0.0])
    def __call__(self, x):
        return {"prediction": self.a * x + self.b}

model = Model(
    Linear(), inputs={"x": Tensor.empty(5)}, targets={"y": Tensor.empty(5)},
    loss=lambda outputs, y: (outputs["prediction"] - y).square().mean(),
)
try:
    model.fit({"x": [-2, -1, 0, 1, 2], "y": [-4, -1, 2, 5, 8]},
              epochs=100, optimizer="sgd", lr=0.1)
    model.save("linear.pgb", include_optimizer=False)
finally:
    model.dispose()

In Node, using a matching Polygrad package:

const { Model } = require('polygrad')
const model = Model.load('linear.pgb')
try {
  const { prediction } = model.forward({x: new Float32Array([3, 4, 5, 6, 7])})
  console.log(Array.from(prediction)) // approximately [11, 14, 17, 20, 23]
} finally {
  model.dispose()
}

See Models for capture, input shapes, minibatches and export options.

Working With Tensors

Creation

from polygrad import Tensor

x = Tensor([1, 2, 3])
a = Tensor.zeros(2, 3)
b = Tensor.ones(2, 3)
c = Tensor.randn(2, 3)
d = Tensor.arange(0, 6).reshape(2, 3)

Math And Indexing

from polygrad import Tensor

x = Tensor.arange(0, 12).reshape(3, 4)
y = x.permute(1, 0).reshape(2, 6)
z = y.relu().sum(axis=1)
picked = x.gather(1, Tensor([[0, 2], [1, 3], [0, 1]], dtype="int32"))

NumPy And Data Updates

import numpy as np
from polygrad import Tensor

arr = np.array([1, 2, 3, 4], dtype=np.float32)
x = Tensor(arr).reshape(2, 2)
print((x * 2 + 1).numpy())

Polygrad tensors are lazy. Use realize() to execute and numpy() when host readback is needed.

import numpy as np
from polygrad import Tensor

x = Tensor.empty((4,), dtype="float32")
x.copy_from(np.array([1, 2, 3, 4], dtype=np.float32))

y = (x * 3 - 1).realize()
print(y.numpy())

x.update_from(np.array([5, 6, 7, 8], dtype=np.float32))
print((x + 1).realize().numpy())

Use copy_from or update_from for repeated loops that should preserve input buffer identity for compiled replay.

Linear Algebra

from polygrad import Tensor

A = Tensor([[4.0, 2.0], [2.0, 5.0]])
b = Tensor([1.0, 3.0])
x = A.solve(b)

print(x.numpy())

Structured linalg methods are portable tensor-composed fallbacks tested against NumPy and Torch. They do not add LAPACK or runtime library dependencies. Current lstsq is solution-only for full-rank tall or square systems.

Training

Use Tensor autograd when you want to control the training loop:

from polygrad import Context, Tensor
from polygrad.nn import Linear, get_parameters
from polygrad.nn.optim import SGD

Tensor.manual_seed(42)
model = Linear(2, 1)
opt = SGD(get_parameters(model), lr=0.01)

with Context(TRAINING=1):
    for _ in range(100):
        opt.zero_grad()
        x = Tensor([[1.0, 2.0], [3.0, 4.0]])
        y = Tensor([[5.0], [11.0]])
        loss = (model(x) - y).square().mean()
        loss.backward()
        opt.step()

print(loss.item())

Layers And Optimizers

from polygrad.nn import Linear, LayerNorm, RMSNorm, Embedding
from polygrad.nn import get_parameters
from polygrad.nn.optim import SGD, Adam, AdamW

Layers include Linear, LayerNorm, LayerNorm2d, RMSNorm, Embedding, Dropout, GroupNorm, Conv1d, Conv2d, ConvTranspose1d, ConvTranspose2d, InstanceNorm, LSTMCell, and BatchNorm. Optimizers include SGD, Adam, and AdamW; each provides step() and zero_grad().

Models

Start with Quickstart: Models for capture, fitting and export. More runnable scripts are in Python examples.

Capture and input rules

• params defaults to named Tensor attributes of the supplied object; functions require explicit closure state. model.summary() inspects metadata without executing or reading weights.
• With loss, construction captures evaluation and training forwards against the same state. The callable runs twice during construction, never during execution. Without a loss, capture uses the current TRAINING value; changing it later does not recapture the graph.
• Set the seed before construction. BatchNorm/RNG updates affect Model-owned state. Authoring Tensor roots and training mode are restored even on failure; arbitrary Python side effects are not. Capture rejects parameter/input assignments and effectful reads.
• Variable-size calls support one bounded leading dimension with fixed trailing dimensions. Save/load preserves the signature; storage currently reserves maximum capacity. Tensor results retain their invocation's values and shape across later calls and Model disposal, but require a live runtime.
• Flat arrays use the signature. Multidimensional NumPy arrays must match the declared shape, not merely its element count. The equivalent JS binding is {data: typedArray, shape: [rows, columns]}.

model.place('CPU:1') and set_device_map(...) accept the same exact native CPU identities. Unsupported accelerator ordinals reject without changing the Model.

Minibatches

• Use model.fit(data, epochs=2, batch_size=32) when every input/target's leading-axis bounds admit 32. Epochs visit samples in input order and return one loss per step. Omit batch_size to repeat one full batch.
• An incomplete final batch rejects before training unless remainder='drop' discards it or remainder='keep' processes the smaller extent permitted by every input's bounds. No padding or shuffling is implicit.
• Tensor datasets are sliced on-device without frontend host readback and may mix with host inputs.

Saving And Loading

model.save() returns bundle bytes; model.save(path) also writes a file. Model.load(bytes_or_path) uses the default runtime, or pass runtime=rt. Each load owns independent state, while tied weights within a Model stay tied. Disposing a Model does not dispose its runtime.

Bundles preserve auxiliary/RNG and optional optimizer state. include_optimizer=False omits optimizer state, not training entrypoints. To resume training, reapply the optimizer configuration after loading. See export products for graph, program and weight exports and version compatibility.

Pretrained models

from polygrad.hf import download_hf, load_hf, generate
import numpy as np

model_path = download_hf("hf-internal-testing/tiny-random-gpt2")
model = load_hf(model_path, max_batch=1, max_seq_len=16)
try:
    tokens = np.array([[1, 2, 3, 4]], dtype=np.int32)
    result = generate(model, tokens, max_new_tokens=2, temperature=1.0, top_k=10)
    print(result)
finally:
    model.dispose()

load_hf supports GPT-2 and Llama configurations with F32/F16/BF16 safetensors. The generate helper above expects GPT-2 input/output names; it is not a Llama generation API. Qwen loading uses the shared C/GGUF path.

models.GPT2(config) and models.Llama(config) build topology without pretrained weights. They reject execution and saving until every required parameter has been written; use a checkpoint loader for a ready-to-run model. Explicit zero weights are valid. Wrong input dtypes are rejected, not silently narrowed.

models.list() (also rt.models.list()) reports construction and checkpoint capabilities; see supported models. Calling models.Qwen3(...) reports that it is import-only and points to Model.from_gguf(...); it does not construct an uninitialized model.

New Qwen3 GGUF imports accept only int32 token input x and return output: model.forward(x=token_ids)['output']. Rotary tables are Model-owned state, included in saved bundles. Older bundles retain their original signatures; use model.entrypoints() to inspect them.

Llama inference is fixed-window, float32, starting at position zero. It supports dense Llama 2/base 3 and text-only 3.x llama3 RoPE scaling, not KV caching, MoE or arbitrary RoPE schemes. The loaders do not supply tokenizers/chat templates.

The C Model GGUF loader converts F32/F16/BF16, I8/I16/I32 and complete Q4_0/Q4_1/Q8_0/Q6_K blocks. Q4_K/Q5_K metadata is readable but weight conversion is unsupported. This is narrower than Python's Tensor-level GGUF decoder; decoding a file does not imply support for its model architecture.

Vision models

models.CLIP, models.ViT, models.DINOv2 and models.DINOv3 use shared C components. Construct from a configuration, write all required weights, or load a Hugging Face directory with Model.from_hf(path). Unloaded models reject execution and saving. Their bundles load in JavaScript with pg.Model.load(bytes).

Model	Inputs	Forward outputs
CLIP	pixel_values, input_ids	image_embeds, text_embeds, logits_per_image, logits_per_text
ViT	pixel_values	last_hidden_state, pooler_output (tanh pooler)
DINOv2 / DINOv3	pixel_values	last_hidden_state, pooler_output (CLS token)

Inputs are already normalized float32 images in [batch, channels, height, width] layout. Use the checkpoint's image processor; the Model does not decode, resize or normalize images. CLIP additionally takes int32, right-padded token IDs with an EOS token. Its encode_image and encode_text entrypoints accept one modality and return normalized embeddings. Both modalities use the configured batch size.

This implementation is fixed-square-resolution, unmasked inference. It supports HF CLIPModel, ViTModel, Dinov2Model and DINOv3ViTModel weights, not classification heads, DINOv2-with-registers, DINOv3 ConvNeXt, training augmentations or positional interpolation to other resolutions. DINOv3 includes register tokens and patch-only 2D RoPE; last_hidden_state retains those prefix tokens. DINOv2 SwiGLU and DINOv3 gated MLP configurations are supported. JSON type tags are clip, vit, dinov2 and dinov3_vit.

For C-built model types and JSON-based models.Sequential / models.Graph, see shared JSON reference. These return the same Model type and use the same training and export APIs. Their configurations support typed inputs, one bounded leading batch dimension, and shared embedding, normalization, RoPE and attention components. For example, {"dtype":"int32","shape":[{"name":"batch","min":1,"max":32},16]} declares token batches of 1 to 32 rows without rebuilding the model.

Devices And Runtimes

Choose a device explicitly or through the environment:

from polygrad import Device, Tensor

x = Tensor.rand(4)

if Device.cuda_available():
    y = (x * 2).to("cuda")
else:
    y = (x * 2).to("cpu")

print(y.numpy())

Environment variables:

DEV=CPU python app.py
POLY_DUMP_KERNELS=1 python app.py
BEAM=4 python app.py

Device selection uses an explicit device=, then POLY_DEV, then DEV. Names are case-insensitive; available native backends include CPU, CUDA, X86 and INTERP. Unsupported device ordinals and multi-device targets are rejected.

Explicit runtimes are optional. Tinygrad-style programs use from polygrad import Tensor, nn and from polygrad.nn.optim import Adam without create() or a Model wrapper. The Python package uses a module-level default C context. Caller-created tensors share that context, so package functions should accept and return Tensor objects rather than copying through NumPy unless readback is required.

For explicit context/device ownership, create a runtime:

import polygrad

rt = polygrad.create(device="cpu")
try:
    layer = rt.nn.Linear(2, 1)
    x = rt.Tensor([[1.0, 2.0]])
    print(layer(x).numpy())
finally:
    rt.dispose()

Use explicit runtimes for isolation, device-specific package wiring, or tests that need independent compiler caches.

Use rt.nn for layers, optimizers and state loaders, rt.models for C-backed model types, and rt.Model for capture/import on that runtime. For example, rt.Model.load("model.pgb") is equivalent to polygrad.Model.load("model.pgb", runtime=rt). Create tensors and layers on the same runtime: mixing owners is rejected even when their devices match; package-level constructors use the default runtime.

Threads

Never use one runtime concurrently from multiple Python threads: native calls release the GIL, and overlapping context mutations can crash the process. The default runtime is process-global. For parallel work, create one runtime per thread, construct through rt.Tensor and rt.Model, and keep its objects in that thread, including cleanup. Separate-runtime CPU, INTERP and X86 execution is tested; concurrent GPU initialization is not validated. Alternatively, serialize all use and cleanup of a shared runtime yourself. Global configuration changes still need coordination. Call collect() and dispose() only when no other thread is using the runtime.

Runtime Inspection

import polygrad

print(polygrad.stats())
print(polygrad.can_run("add", shape=[1024]))

can_run(...) is conservative. For some compound op/shape queries it raises when support cannot be proven statically.

JIT And Compile

@jit follows tinygrad's raw Tensor JIT behavior. The first call runs normally, the second call captures realized schedules, and later calls replay those schedules with current input buffers.

from polygrad import Tensor, jit

@jit
def step(x):
    return (x + 1).realize()

print(step(Tensor([1, 2, 3])).numpy())  # normal run
print(step(Tensor([4, 5, 6])).numpy())  # capture
print(step(Tensor([7, 8, 9])).numpy())  # replay

For embedding loops, compile(...) performs the same warmup and capture up front and exposes an explicit callable:

from polygrad import Tensor, compile

def step(x):
    return (x + 1).realize()

compiled = compile(step, [Tensor([1, 2, 3]).realize()])
out = compiled.run([Tensor([7, 8, 9]).realize()])

print(out.numpy())
print(compiled.stats())
compiled.dispose()

polygrad.stats() exposes shared C runtime counters. The tinygrad-compatible top-level GlobalCounters provides global_ops, global_mem, time_sum_s, kernel_count, mem_used, mem_used_per_device, and reset(); reset keeps live allocation bytes. Jit.stats() and CompiledCallable.stats() expose wrapper-level capture and replay counters.

polygrad.can_run(op, dtype="float32", shape=..., shapes=..., device="auto") is an advisory backend capability probe.

Capture requires at least one Tensor argument. Host reads must stay outside the captured function; otherwise their values would be fixed during capture. Compiler settings such as BEAM and NOOPT are documented in the shared runtime reference.

Reusing Input Buffers

import numpy as np
from polygrad import Tensor, compile

x = Tensor(np.array([1, 2, 3], dtype=np.float32)).realize()
f = compile(lambda x: x.square().sum().realize(), [x])

print(f.run([x]).item())
x.copy_from(np.array([4, 5, 6], dtype=np.float32))
print(f.run([x]).item())
f.dispose()

Custom Kernels

Tensor.custom_kernel(...) mirrors tinygrad's alpha custom-kernel shape. The kernel function receives placeholder UOps and returns a SINK body. Polygrad wraps the body in CALL, returns AFTER(...) tensors, and keeps execution in the normal schedule and runtime caches.

from polygrad import Tensor
from polygrad.uop.ops import KernelInfo, UOp

def add_kernel(out, a, b):
    out, a, b = out.flatten(), a.flatten(), b.flatten()
    i = UOp.range(out.ctx, out.numel(), 0)
    return out[i].store(a[i] + b[i]).end(i).sink(
        arg=KernelInfo(name="custom_add_4")
    )

out = Tensor.empty((4,), dtype="float32")
y = out.custom_kernel(
    Tensor([1.0, 2.0, 3.0, 4.0]),
    Tensor([10.0, 20.0, 30.0, 40.0]),
    fxn=add_kernel,
)[0]
print(y.numpy())  # [11. 22. 33. 44.]

KernelInfo is required to mark the executable kernel boundary. Without it, the body can leave the output unchanged, as in pinned tinygrad. Match each stored value's dtype to its destination, using an explicit UOp cast when needed. Current C and Wasm renderers reject mismatched vector stores with vector STORE dtype mismatch; cast the value to the destination dtype. Scalar C stores can convert numerically, but that is not a portable kernel contract. INTERP follows Tinygrad's Python memoryview conversion rules. Cast explicitly on every backend.

Package Integration

Python packages should accept caller-created Polygrad tensors and return Polygrad tensors:

from polygrad import Tensor

def normalize(x: Tensor) -> Tensor:
    mean = x.mean(axis=-1, keepdim=True)
    scale = (x - mean).square().mean(axis=-1, keepdim=True).sqrt()
    return (x - mean) / scale

This keeps execution in the caller's Polygrad context and avoids unnecessary NumPy readback.

API Overview

Area	Main APIs
Runtime	polygrad.create, polygrad.stats, polygrad.can_run, Device
Models	Model, from_callable, from_tensors, fit, forward, save, load, summary, dispose; polygrad.models factories
Tensor creation	Tensor(data), zeros, ones, full, rand, randn, randint, arange, linspace, eye, empty
Tensor math	+, -, *, /, **, exp, log, sqrt, abs, sin, cos, tanh, sigmoid, relu, gelu, silu, softmax
Reductions	sum, mean, max, min, argmax, sort, argsort, topk, var, std
Movement/indexing	reshape, view, permute, transpose, expand, squeeze, unsqueeze, flatten, shrink, pad, flip, repeat, gather, take_along_axis, cat, stack, split, chunk
Linalg	matmul, dot, linear, qr, triangular_solve, solve_triangular, cholesky, cholesky_solve, solve, lstsq
Data/readback	realize, numpy, item, tolist, copy_from, update_from, to, cpu, cuda, detach, clone
Compilation	jit, compile, Tensor.custom_kernel
Neural nets	polygrad.nn layers, SGD, Adam, AdamW, get_parameters, get_state_dict

Troubleshooting

• unknown type name '__fp16': float16 CPU kernels require Clang. Install it and select CC=clang; check that CC is not forcing GCC.
• args mismatch in jit: the call must match the traced input shapes, dtypes and devices. An int32 input cannot replace a float32 sample.
• Bundle ABI/format mismatch: use matching producer/consumer Polygrad versions. See bundle compatibility; do not edit artifact version fields to bypass validation.
• CPU compilation cannot find a compiler: install clang (recommended) or GCC, or try DEV=X86 on a supported x86 machine or DEV=INTERP for interpreted execution.
• "TRAINING must be enabled": wrap optimizer steps in with Context(TRAINING=1):, as in Training.
• Operands belong to different runtimes: create layers and Tensors through the same runtime (rt.nn, rt.Tensor), or use the default API consistently.
• Memory grows across many different graphs: at an idle boundary, use rt.clear_schedule_cache() and rt.collect(). Live owners still retain their resources; later execution rebuilds cleared schedules. This is a memory-pressure tool, not a per-step requirement.

For contributor checks, see tests.

License

MIT