Polygrad
Polygrad is a tensor library and JIT compiler built around a C11 port of tinygrad. It provides Python and JavaScript APIs for automatic differentiation, neural networks and model export. Native applications can use the C API directly.
- Python guide: installation, Tensor/NumPy usage, training, Models and JIT.
- JavaScript guide: Node, browser/WebGPU, typed arrays, training and Models.
Install
Python:
pip install polygrad
JavaScript:
npm install polygrad
Python requires Linux and Python 3.9+; Node.js requires Node 18+. See the Python installation guide and JavaScript installation guide for compiler requirements and native/Wasm setup.
Quickstart: Tensors
Python:
from polygrad import Tensor
x = Tensor.rand(3, 4)
w = Tensor.rand(4, 5)
y = (x @ w).softmax(-1)
print(y.numpy())
Node.js:
const { Tensor } = require('polygrad')
const x = Tensor.rand(3, 4)
const w = Tensor.rand(4, 5)
const y = x.dot(w).softmax(-1)
console.log(y.toArray())
Browser (npm package with a browser bundler):
import { Tensor } from 'polygrad'
const y = new Tensor([1, 2, 3]).mul(2).add(1)
console.log(y.toArray())
Polygrad follows tinygrad’s graph-based execution model:
- Build a graph: Tensor operations create UOps describing the computation, rather than immediately running a kernel for each operation.
- Lower and render: when results are needed, the scheduler forms kernels; the compiler lowers their operations and renders backend-specific code.
- Execute: run on native CPU or CUDA, or use WebAssembly and WebGPU in the browser. Direct x86 and interpreter backends are also available.
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()
Load in Node without the Python class:
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()
}
Polygrad adds three pieces around tinygrad’s compiler core:
- Model and helpers: capture Tensor computations or construct supported architectures, with named inputs, outputs and state.
- Logical and physical graphs: retain a portable logical graph separately from the device-specific graph used for execution.
- Portable bundles: save the logical graph and weights together. Another frontend can load the bundle and compile it for its backend without the original model class.
Devices And Runtimes
Explicit runtimes and Models are optional for ordinary Tensor programs.
Python Tensor API JavaScript Tensor API C / native package
| | |
+------------------------+-------------------------+
|
Polygrad C11 runtime
|
+----------------+-------+-------+----------------+
| | | |
CPU/x86 CUDA/HIP WASM WebGPU
| Situation | Use |
|---|---|
| Python research or scripts | polygrad from PyPI |
| Node.js product code | polygrad from npm with the default Tensor API or polygrad.create() |
| Browser CPU/WASM | polygrad with core: "wasm" |
| Browser GPU | polygrad with core: "wasm", device: "webgpu" |
| Native embedding | C ABI and libpolygrad |
| Package integration | create one runtime and pass it into the package |
Browser WebGPU execution and readback require async methods. HIP is implemented but excluded from the release-validation matrix.
Choose a device through Tensor/Runtime options or, in native applications,
POLY_DEV / DEV. Use the default runtime for ordinary programs; create an
explicit runtime for isolation and construct related objects through it.
See Python runtimes,
JS runtimes, and the
shared settings and ownership reference.
Current Limits
- Python installation currently targets Linux.
- Polygrad tracks tinygrad 0.14.0, but is not a drop-in replacement or a distributed-training framework. Genuine multi-GPU execution is unsupported.
- Do not use one Python runtime concurrently from multiple threads, including the default runtime. Overlapping calls can crash the process; use a runtime per thread or serialize all use and cleanup.
- Integer division by zero at valid coordinates can terminate native Python or Node. Validate divisors supplied by users.
- Structured linear algebra uses portable Tensor compositions, not vendor BLAS/LAPACK kernels.
See compatibility for backend and API gaps, and the Python or JS troubleshooting sections for installation and runtime errors.
Reference
Guides By Topic
- JIT: Python, JavaScript.
- Custom kernels: Python, JavaScript.
- Package integration: Python, JavaScript.
Supported Models
| Model type | JSON construction | HF import | GGUF import | Generation |
|---|---|---|---|---|
| MLP, TabM, NAM, Sequential, Graph | Yes | No | No | No |
| GPT2, DistilGPT2 | Yes | Yes | Yes | Legacy Python helper (uncached) |
| Llama | Yes | Yes | No | C-backed Transformer in Python/JS |
| Qwen3 | No | No | Yes | C-backed Transformer in Python/JS |
| Qwen35Text | Yes | Text backbone only | No | Hidden states with optional caches; no token generation |
| Qwen35 | Yes | Text and packed images | No | Cached logits/sampling; external image preprocessing and fixed image grids |
| Whisper | Yes | Yes | No | Internal caches and fused token sampling; external audio preprocessing |
| CLIP, ViT, DINOv2, DINOv3 | Yes | Yes | No | No |
| EmbeddingGemma2 (text/image/audio), EmbeddingGemma2Text | Yes | Yes | No | No |
Checkpoint-required types must have all weights loaded before execution or
export. models.list() reports construction/import capabilities. Cached Llama/Qwen3 generation
requires cache_capacity / cacheCapacity; specialize an imported Model with
Transformer.from_model() / Transformer.fromModel().
Tokenization is separate from model loading. Python Tokenizer.from_json and
JavaScript Tokenizer.fromJSON use the same C tokenizer and accept validated
byte-level BPE pipelines; unsupported normalizers and tokenizer algorithms
reject. Explicit Python strict=False / JS {strict:false} (C:
poly_tokenizer_from_json_ex) can skip NFC with a warning, at the cost of
potentially different token IDs. Both frontends also expose GGUF tokenization.
See tokenizer support.
HF import accepts safetensors and a restricted subset of PyTorch ZIP state
dictionaries (.bin/.pt). It reads weights without executing Python code;
it does not load pickled models or add support for an unknown architecture.
Python and JavaScript share checkpoint file selection in C; network access
stays in the frontends. See the Python
and JavaScript guides for selection,
supported formats and limits.
EmbeddingGemma 2 produces text, image and audio embeddings using external Hugging Face preprocessing; see the Python and JavaScript guides for inputs and modality selection.
Configuration-driven Models
Sequential and Graph use the same JSON definitions in Python and JavaScript.
Named factories accept untagged configs; Model(config) / new Model(config)
requires format: "poly.modeldef@1" and a registered type.
models.list() reports which types support construction or checkpoint import.
These factories return ordinary Models, not architecture-specific subclasses.
The JavaScript snippets below use an existing runtime pg and parsed JSON config.
JSON examples, components and construction limits (Python and JavaScript)
For example, save this as network.json:
{
"input": {"name": "x", "shape": [1, 4], "dtype": "float32"},
"layers": [
{"name": "first", "type": "linear", "out_features": 2, "activation": "relu"},
{"name": "second", "type": "linear", "out_features": 3, "activation": "relu"},
{"name": "head", "type": "linear", "out_features": 4}
],
"output": "prediction",
"seed": 42
}
from pathlib import Path
from polygrad import models
model = models.Sequential(Path("network.json").read_text())
# Optional runtime=rt uses an explicit Runtime instead of the default context.
const model = pg.models.Sequential(config) // Node/native or synchronous Wasm
// WebGPU: await pg.models.SequentialAsync(config)
// Graph and GraphAsync accept the connected form described below.
models.DistilGPT2(config) (JS: pg.models.DistilGPT2(config)) is a six-block
GPT-2 preset using the same builder and checkpoint names. Config fields can
override defaults. Like GPT-2, it requires weights before execution or export;
HF DistilGPT-2 checkpoints still declare model_type: "gpt2" and load normally.
The initial component catalogue is deliberately bounded:
| Component | Configuration |
|---|---|
linear |
out_features; optional bias (default true), activation (default none) |
embedding |
Integer indices; vocab_size, embed_dim |
layernorm, rmsnorm |
Normalize the fixed last dimension; eps (defaults 1e-5/1e-6), affine (default true) |
rope |
Split-half rotation of [batch,heads,sequence,head_dim]; fixed sequence and even head width; theta (default 10000) |
attention |
Query, key, value, optional mask; is_causal, enable_gqa (default false); no dropout |
relu, sigmoid, tanh, silu, gelu |
One input |
identity, square, neg, exp, log, softplus |
One input; softplus uses beta 1 |
add, sub, mul, div |
Two inputs, existing Tensor broadcasting |
matmul |
Two inputs, existing Tensor dot semantics |
maximum, lt, eq |
Two inputs, existing Tensor broadcasting/comparisons |
where |
Condition, true value, false value |
arange |
No inputs; integer start, stop, positive step, concrete dtype; nonempty bounded range |
const_like |
One floating input; finite value, with the input’s shape and dtype |
gradient |
Two inputs: floating scalar loss and floating target; constructs its derivative graph |
sum, mean, max |
Optional nonempty axes and boolean keepdim; otherwise reduce all axes |
reshape |
Positive shape, optionally the same bounded leading dimension; unchanged symbolic element count |
permute |
axes: a permutation of all input axes |
cast |
Explicit destination dtype |
repeat |
Positive integer count; body is one unnamed component or a named layer list |
Graph configurations replace input/layers/output with:
inputs: name maps to{shape, dtype, role?}; role isinputortarget. An empty mapping is allowed for state-only entrypoints.- Optional
state: name maps to{shape, dtype}. Fixed-shape AUX buffers start at zero, persist across calls, and are included in checkpoints. - Optional
controls: name maps to integer{min, max}bounds. Controls are scalar values available to nodes, not shape declarations. nodes: ordered{name, type, inputs: [earlier_value_names], ...}records.outputs: output name maps to input or node name.- Optional
entrypoints:{name, inputs, outputs, objective?, controls?, writes?}records. Without them,forwardexposes all declared inputs and outputs. An explicit scalar objective enables the existing Model training path.controlslists declared control names;writesmaps state names to node/value names of exactly the same shape and dtype. Writes use pre-update values and commit together on a successful call (execution failures are not rolled back). Writing entrypoints require distinct output names. At least one output is required, even when it is only a scalar status.
Stateful loops can call model.call('step', {}, controls={...}, readback=False)
in Python or await model.callAsync('step', {}, {controls: {...}, readback: false})
in JS. This returns an empty mapping without collecting outputs; use
read_buffer / readBufferAsync for occasional snapshots. It does not make an
in-flight GPU operation cancellable or relax runtime serialization.
See the state/control example.
Both model types accept modules, a table of named leaf-component configurations.
Use {name, call: "shared", inputs: [...]} instead of type in a Graph node
(omit inputs in Sequential). Here shared is the key of a declaration in
modules, not a reserved keyword. Repeated calls reuse the declared component’s
parameters; ordinary Repeat bodies create fresh parameters. Module declarations
are construction components, not device-placement cuts or runtime submodels.
See the shared-layer Graph configuration.
Parameter names are layers.<name>.weight/bias, nodes.<name>.weight/bias, or
modules.<name>.weight/bias; Repeat inserts zero-based indices. Linear weights
and embedding weights use the existing seed/name-keyed C model Kaiming uniform
initializer; biases are zero. Normalization weights start at one. Parameters are
float32; Tensor promotion rules apply to other input dtypes. RoPE owns deterministic
freqs_cos/freqs_sin AUX tables under its component name. This does not promise
Keras/Tinygrad initial-weight equivalence or automatic conversion between split-half
and interleaved checkpoint layouts.
Inputs require an explicit concrete scalar dtype supported by the selected backend,
rank <=8 and positive dimensions.
A leading dimension may be {"name":"batch","min":1,"max":32}. All bounded
declarations in one definition must use that same name and bounds; trailing
dimensions remain fixed. Calls resolve concrete extents through the normal Model
input contract, and reductions use the invocation extent, not the declared maximum.
See the typed, bounded transformer-component fixture.
Names are ASCII identifiers of 1-63 characters. Configuration limits
are 1 MiB JSON, nesting 32, 16,384 JSON values, 1,024 expanded component calls,
construction depth 16, and 64 inputs/outputs/entrypoints. Expanded paths are at
most 191 bytes. Named storage totals at most 16,777,216 elements; node shapes,
linear contractions, embedding selectors and attention scores have the same limit,
evaluated at maximum extents.
These are construction limits, not a bound on compiler/backend peak memory.
Optional format: "poly.modeldef@1" and type: "sequential"/"graph" tags are
checked when present. The selected factory already identifies the model type.
Unknown fields, duplicate keys/names, forward/cyclic references, unused shared
components and incompatible shapes fail. There are no config expressions,
recursive modules, runtime loops, data-dependent shapes or Keras JSON compatibility.
Configuration describes construction, not checkpoint state: save/load the
result through existing Model bundle or graph/weights APIs.
ONNX Import
Python Model.from_onnx and JavaScript Model.fromONNX share a C loader that
lowers into ordinary Tensor/UOp graphs and seals a portable Model. No ONNX
runtime or Python dependency is needed at execution. See the
Python and JavaScript calls.
Call the imported Model with its ONNX input names, or save it as a Polygrad bundle to run elsewhere without redefining the model. Import is currently inference-only; unsupported operators and shapes produce an import error.
Export Products
Use save() / load() for a portable graph-and-weights bundle.
Newer releases load older portable models; a model saved by a newer release
may require that release or later. Check the release notes before exchanging
models between different package versions.
Cached models save their program and weights, not conversation history.
For separate artifacts:
export_ir()/exportIR()returns the portable model graph. Import it with a new placement policy when the target device layout may change.export_program()/exportProgram()returns the currently compiled, device-bound program. It starts without rebuilding the model graph, but requires a compatible Polygrad binary interface and backend/device.export_weights()/exportWeights()returns named safetensors state. Pass it separately to either import path when the model has parameters/state.
Python:
program = model.export_program()
weights = model.export_weights()
fast_model = Model.from_program(program, weights)
result = fast_model.call("forward", {"x": input_array})
JavaScript (native/Wasm synchronous runtimes):
const program = model.exportProgram()
const weights = model.exportWeights()
const fastModel = Model.fromProgram(program, weights)
const result = fastModel.call('forward', { x: inputArray })
WebGPU startup is asynchronous, so use await model.exportProgramAsync() and
await Model.fromProgramAsync(program, weights). A bound-program Model
is inference/call-only: it has no portable logical graph and cannot be
placed on another device, trained, differentiated, or exported as a portable
graph. Use a regular Model bundle when portability is needed.
Settings
Compiler tuning
Python Context(NOOPT=1) and JS runtime.noopt = 1 disable automatic
kernel scheduling heuristics, not the compiler’s required lowering passes or
explicit BEAM search. The setting participates in program-cache keys and does
not change already compiled JIT captures. It is shared by native contexts;
each Wasm module has its own setting. Restore it after a scoped experiment.
Python Context(BEAM=2) and JS runtime.beam = 2 select kernel beam-search
width through the same C policy. Zero disables search. Search compiles and
times candidates on the selected backend; WebGPU timing requires the adapter’s
optional timestamp-query feature. Without it, candidates are not timed and
the original kernel is kept. Change JS settings while idle and restore them
after use; existing compiled JIT captures are unchanged.
Search compares its winner with the default schedule at full size and caches
the selection. Python Context(BEAM_COMPARE=0) or JS runtime.beamCompare = 0
disables this comparison to use the pinned tinygrad search policy.
Wasm uses a no-pack SIMD128 UOp matmul kernel for supported fixed FP32 shapes.
It uses strict multiply/add. Contiguous K rows use four partial sums per output,
combined as (s0+s2)+(s1+s3); contiguous output columns use serial K accumulation.
The striped order changes rounding, especially under cancellation; it is not
guaranteed to be more accurate than serial accumulation. Relaxed multiply-add
is not used. Unsupported shapes retain the existing fallback path.
create({ kernels: false }) or POLY_KERNELS=0 disables kernel selection.
Fallbacks and disabled kernels use shared optimization and ordinary Wasm rendering;
the renderer does not replace matmuls or row reductions with specialized kernels.
Runtime(kernels=True) in Python or create({ kernels: true }) in JavaScript
additionally enables experimental AVX2/FMA CPU and tiled FP32 WebGPU kernels.
These remain opt-in and are not guaranteed speedups. POLY_KERNELS=1 enables
all providers. Policy is fixed at first compilation; portable bundles keep
the original operations.
The explicit opt-in policy also optimizes FP32 interval sums expressed as
((arange(N) >= lo) & (arange(N) < hi)).where(values, 0).sum(1), with
int32 bounds [S,1,1] and values [1,N,W]. CPU, X86, interpreter, Wasm,
CUDA and WebGPU execute direct interval loops instead of scanning all N values
for every segment. Bounds clip to [0,N]; empty or reversed intervals return
zero. Overlapping intervals are allowed. This is a compiler optimization of
existing Tensor operations, not a new sparse Tensor API. Authoring, gradients
and portable bundles keep those operations; unsupported patterns fall back.
Serial summation can round differently from the default reduction tree.
POLY_KERNELS=1 sets the default for newly created runtimes; an explicit
kernels option overrides it. Changing the environment does not change an
existing runtime. The CPU path also materializes bounded softmax tables to
avoid repeated exponentials. Custom kernels can change rounding; benchmark
your workload before enabling them.
Python Context(IGNORE_BEAM_CACHE=1) or JS runtime.ignoreBeamCache = 1
bypasses saved search results. BEAM_TIMEOUT_SEC bounds candidate compilation
(default 10 seconds; 0 disables the budget). Native compiler subprocesses are
stopped on expiry; in-process driver calls are checked when they return.
BEAM_STRICT_MODE=1 also propagates native compiler process failures instead
of rejecting those candidates. Ordinary candidate-limit and device compiler
rejections remain recoverable, matching the pinned search policy.
For native search diagnostics, BEAM_DEBUG=1 prints the input graph and final
options; BEAM_DEBUG=2 also prints candidate compilation and execution times.
These diagnostics are disabled by default and use the existing C graph printer.
Default dtypes
Default storage dtypes use the same C-owned policy: Python
Context(DEFAULT_FLOAT='float64', DEFAULT_INT='int64'), or JS
runtime.defaultFloat / runtime.defaultInt. Change JS settings while idle
and restore them after use. Explicit dtypes and typed-array inputs are preserved;
backend dtype capabilities still apply. These settings participate in compiler
cache keys. They do not change fixed accumulation/RNG compute widths or the
bounds-based int32/int64 lowering of weak integer indices.
DISK storage
Native file-backed tensors use C-owned DISK:<path> storage. Python
Tensor(pathlib.Path(...)) passes the path to C for memory mapping;
.to('DISK:<path>') copies to that exact file, preserving filename case.
Tensor.empty(..., device='DISK:<path>') constructs lazy storage with that
identity; it does not open or initialize the file. JS uses
Tensor.empty(shape, {device: 'disk:<path>'}).
DISK is not a compute backend. Wasm does not implement native file mapping:
load bytes in JavaScript and pass typed arrays through the host-buffer API.
JS does not currently provide Python’s Tensor path constructor.
Environment variables
DEV=CPU # also CUDA, HIP, X86, INTERP, CPU:X86
POLY_DEV=CPU # Polygrad-specific override of DEV
POLY_LIB=/path/to/libpolygrad.so
POLY_CORE=native # or wasm
POLY_DUMP_KERNELS=1
BEAM=4
Explicit device choices override POLY_DEV, which overrides Tinygrad’s DEV.
Both environment names accept case-insensitive single backend names and
CPU:X86. Unsupported renderer, architecture, interface, ordinal-like and
multi-target forms fail rather than silently selecting another backend.
Not all Tinygrad target-selection forms are supported; multi-GPU execution
is also unsupported.
DEV=CUDA:1 is a renderer request in Tinygrad’s target grammar, not GPU1.
Device strings passed directly to Tensor/Runtime retain their existing rules.
POLY_DEBUG overrides DEBUG, including POLY_DEBUG=0 to suppress inherited
verbosity. BEAM and NOOPT retain their Tinygrad names; explicit setters
override their initial environment values. Native runtimes share compiler
policy; each Wasm module owns its policy. Node initializes Wasm compiler
settings once per module. Browsers use runtime properties, not process env.
POLY_LIB is Python’s explicit native-library path; an invalid Python path
fails without falling back to another installation. Node loads its packaged
addon or Wasm via POLY_CORE. The obsolete POLY_DEVICE and POLYGRAD_LIB names are not recognized.
Logical Graphs And Placement
default Tensor construction
-> eager tinygrad-shaped physical tensor graph
-> LINEAR schedule graph with CALLs
-> PROGRAM/SOURCE/BINARY plus runtime runner
portable logical program + named state + explicit policy
-> place from scratch
-> replacement complete physical tensor graph
The key invariant is that logical roots never drive default execution or alter the Tinygrad-shaped physical graph. Logical lifetime is configurable:
until_realize(default) keeps a producer until that Tensor is materialized, then retains only an exact device-free BUFFER/movement resource;alwayskeeps the full producer graph for portable export or later placement;neverbuilds only the physical graph and cannot provide portable export.
Use until_realize for repeated eager training or assignment loops. always
also retains previous updates, so their graph and differentiation cost grow
with the loop. Scope always to portable graph authoring/capture instead;
captured Models own their programs independently of the source Tensors.
If any operand has no logical root, a new pure result remains physical-only even after the ambient context returns to another policy. This never changes the eager physical graph.
Set the context default with POLY_LOGICAL=2|1|0. Python also supports
Context(LOGICAL=...), Runtime(logical=...), runtime.logical(...),
Tensor(..., logical=...), and tensor.preserve_logical(). JavaScript supports
createRuntime({logical: ...}), runtime.withLogical(...),
new Tensor(data, {logical: ...}), and tensor.preserveLogical().
Model capture owns its logical entrypoints immediately, so later realization
or retirement of the source Tensor cannot invalidate the portable program.
Default Tensor/JIT execution builds and schedules the physical graph eagerly;
it does not invoke placement. Retained Model graphs may instead request an
explicit policy. Device maps use exact named module cuts:
from polygrad import Model, Tensor
x = Tensor.empty(2)
h = x + 3
y = h * 2
model = Model.from_tensors(
inputs={"x": x}, outputs={"y": y},
modules=[
{"name": "stem", "inputs": [x], "output": h},
{"name": "head", "inputs": [h], "output": y},
],
)
model.set_device_map({"stem": "CPU", "head": "CPU:1"})
assert model.forward(x=[1.0, 2.0])["y"].tolist() == [8.0, 10.0]
model.dispose()
The policy places the aggregate retained graph once, keeps module state with
its module, and inserts explicit COPY nodes at cross-device cuts. CPU sibling
identities are executable today; nonzero CUDA/HIP identities, automatic
sharding, pipeline schedules, offload, and VRAM planning remain future work.
PGIR preserves the exact named module boundaries but not their device
assignments, so an imported program can be placed under a new map.
Runtime Ownership
Runtime owners, schedule-cache growth and reclamation
A Runtime owns compiler caches, buffer residency and backend handles. Ordinary imports use the default runtime; explicit runtimes provide isolation. Use one owner for related Tensors, layers and Models. See the Python and JavaScript guides.
When realization reports an error, Tensor roots return to their previous expressions. Preparation failures before any dispatch leave existing values intact and allow retry. After partial execution, storage written by attempted calls is unusable; reads and reuse fail rather than return partial results. Recreate affected state from a checkpoint. This does not roll back writes or recover a lost GPU device.
The schedule cache has no automatic eviction or size cap. Each entry keeps a
content digest and its schedule, not the source graph. Graphs with host callbacks
or auxiliary objects still use retained pointer keys. Long-lived runtimes
producing many distinct graphs can accumulate schedules; collect() alone does
not evict them. Per-node digest caching also adds bookkeeping memory.
Symbolic bindings that reuse a schedule do not necessarily add cache entries.
At an idle boundary, use Python pg.clear_schedule_cache() or
runtime.clear_schedule_cache(), JavaScript runtime.clearScheduleCache(),
or C poly_schedule_cache_clear(ctx) from polygrad.h. Then call
collect() / poly_ctx_collect(ctx) to reclaim resources without other owners.
Storage reclamation is deferred to safe points, including the next Tensor
readback after a storage-owning root is retired. Ordinary JIT/readback does not
scan retained graphs merely because a temporary wrapper was released.
Await pending browser operations before clearing; raw C callers must serialize
access and finish queued device work. Independently owned Model/JIT executions
remain usable. Use clearing at workload boundaries or under memory pressure,
not after every operation. New schedules may need rebuilding, but compiled
kernels can still be reused from other caches; recompilation is not inevitable.
Other caches are untouched, so clearing does not promise that all runtime
memory is returned.
Tinygrad Relationship
Polygrad ports tinygrad’s compiler concepts to C: UOps, rewrites, LINEAR schedules,
PROGRAMs, JIT capture/replay and backend lowering. The pinned reference is
tinygrad v0.14.0, commit
6f87158d77f66a36d5f8bbe915170b24e2acabe8.
Compatibility limits and intentional differences
The public Tensor APIs do not yet provide SVD/Newton-Schulz,
nonzero/masked_select (including fixed-size variants), or borrowed-pointer from_blob
construction. Tinygrad’s PYTHON backend name is not an alias for Polygrad’s
INTERP backend. These are compatibility limits, not passing upstream tests;
the reviewed Tensor,
operation and
NN/optimizer baselines
keep nonpassing cases explicit, including unsupported Python compiler-private
helpers.
Tensor shapes support at most 64 dimensions. Portable Model inputs and outputs support at most 8 dimensions.
WebGPU does not guarantee NaN truthiness or preservation through clipping under
WGSL finite-math rules. Unsupported features and reviewed divergences are listed
in the register.
Entries labeled kind: upstream_backport identify fixes from newer tinygrad
commits, separately from Polygrad-specific policies. Their exact upstream_commits
must be included in a future pin, with affected parity and regression tests passing,
before closing the allowance. A version bump alone does not close it; the regression
tests remain afterward.
The main intentional differences are:
| Area | Polygrad difference |
|---|---|
| Core runtime | Compiler state, buffers, caches, and backend runners live in PolyCtx inside a C library |
| Frontends | Python and JavaScript are wrappers over the same C core rather than separate runtimes |
| WASM/browser | Browser execution uses the unified C/WASM runtime path, with WebGPU orchestrated from the C backend |
| Model tooling | Models provide named inputs, outputs and state, training helpers, and portable save/load |
| Custom kernels | Public custom kernels lower into UOp CALL bodies and still run through normal scheduling and runtime caches |
| WebGPU int64 | WGSL has no native 64-bit integers, so renderer lowering uses two 32-bit lanes while C, CUDA, HIP, WASM, and x86 retain native int64; unlike pinned tinygrad, valid dynamic/uint32 shift counts and signed right shift are handled rather than crashing or changing sign semantics |
| WebGPU narrow integers | Preserve 8/16-bit integer overflow and cast behavior, correcting lost narrowing in the pinned WGSL renderer |
| WebGPU scheduling | Use different matrix-vector tiles for small outputs and longer reductions; other shapes and explicit tuning settings retain the reference policy |
| CPU scheduling | Use bounded maximum sizes for native CPU tiling and threading; small runtime batches avoid worker wake-ups |
| Assignment | Assigning realized data to an unrealized expression produces readable, independent destination storage; the pinned reference can lose that assignment |
Embedding The C Core
See the C examples for Model construction, training and export.
Headers, context ownership and retained roots
The C core owns graph construction, scheduling, placement, runtime caches, and backend dispatch. Frontends are thin wrappers over the same concepts.
For C embedding, polygrad.h is the convenience umbrella. Its domain headers
also work independently: core.h owns shared types and context controls,
tensor.h owns poly_tensor_* operations, and uop/ and mixin/ headers own
poly_uop_* graph construction and composition. Include model.h or nn/nn.h
when using those APIs. frontend.h is separate: it adapts dtype IDs, flattened
arguments and opaque handles for language bindings. Core implementations do
not call these adapters. Use Model-owned declarations or
poly_model_from_bindings for named state, and poly_model_* helpers for layers.
C Model construction uses poly_model_from_config(ctx, type, json, len, device,
&error) from models/models.h. A non-NULL context is borrowed and must outlive
the Model; explicit NULL requests an owned context. Typed constructors include
poly_mlp_into, poly_gpt2_into and poly_qwen3_into. Checkpoint import uses
poly_hf_load / poly_gguf_load, or _into variants for a caller-owned context.
Low-level UOp.variable bounds retain integer, floating-point and boolean
endpoints independently of the variable dtype. C takes scalar PolyArg values;
Python accepts int/float/bool; JavaScript uses pg.uop.variable(...), with
BigInt for exact wide integers. NaN, reversed and nonnumeric bounds are rejected.
Typed endpoints can exceed the runtime’s signed64
variable-binding domain; metadata support does not imply executable bindings.
Raw C PolyUOp * results are borrowed. Retain roots crossing execution or
collection boundaries with poly_uop_retain(), and release them with
poly_uop_release(). Allocation-only loops that never execute must call
poly_ctx_collect(); poly_ctx_stats() does not collect.
Versioning
Python and JavaScript share the same core release line; wrapper-only fixes may increment only the affected package. Additive APIs and internal fixes preserve existing binary interfaces. Binary and saved-model compatibility are checked separately from package versions; incompatible changes are identified in the release notes.
Building From Source
make
POLY_LIB=$PWD/build/libpolygrad.so PYTHONPATH=py python - <<'PY'
from polygrad import Tensor
print((Tensor([1, 2, 3]) * 2 + 1).numpy())
PY
cd js
npm install
node - <<'JS'
const { Tensor, disposeDefault } = require('.')
const y = new Tensor([1, 2, 3]).mul(2).add(1)
console.log(y.toArray())
disposeDefault()
JS
Build the standalone browser bundles from the repository root:
make wasm-pkg
cd js
npm run build:browser
Package builds regenerate binding glue from C headers and
scripts/binding_rules.json, requiring Python, Clang and clang-format-14.
Published packages include the glue; installing them does not run the generator.
For portable Python wheels, run make build-py-manylinux from the host with
Apptainer installed. It builds from the staged release sdist, audits and tests
CPython 3.9-3.13 wheels, then stages them under build/release/<version>/wheels/.
It does not publish packages; logs are under build/manylinux/run.*/build.log.
Tests
make test-readme checks links and runs the local Python/Node examples,
including setup for multi-block workflows. Shell blocks are syntax-checked;
JSON is parsed. Browser blocks are explicitly excluded, and the HF download
example requires README_NETWORK=1. These exclusions are reported as skips.
Developer checks, release acceptance and benchmarks
Functional checks
make test
make test-parity
make test-py
make test-js-native
TMPDIR=$PWD/temp/cc_tmp EM_CACHE=$PWD/temp/emscripten-cache make test-js-wasm
make test-browser # use an existing working DISPLAY
make test-browser-matrix
test-all runs the functional backend/frontend matrix. test-browser covers
real browser WebGPU; test-browser-matrix adds other browser executables:
BROWSER_MATRIX="chromium,firefox,old-chrome=chromium@/path/to/chrome" make test-browser-matrix
BROWSER_MATRIX_DEVICES="auto,interp" make test-browser-matrix
Browser specs are chromium, firefox, webkit, chrome, or
label=engine@/absolute/path.
Browser performance comparisons
For Wasm performance changes, compare against a frozen browser bundle built in a separate checkout, using the same toolchain:
make bench-wasm-regression BENCH_WASM_BASELINE=temp/old-checkout/js/dist/polygrad.sync.js
This first qualifies the measuring method on identical copies of the frozen
baseline, then compares correctness-checked Polygrad workloads in balanced
fresh-page pairs. It does not run jax-js. Reports under temp/performance/
retain samples, bundle/kernel hashes and browser identity; the baseline is never
overwritten. The default slowdown budget is 10%, adjustable with
BENCH_WASM_MAX_SLOWDOWN. A completed comparison exits 1 on regression or 2
when inconclusive; failed qualification prevents comparison. Setup and
correctness failures also fail the command. Do not repeat runs until one passes.
Compare
both the previous release and a known-fast historical baseline when investigating
a regression; passing against the immediate parent alone is not recovery.
The same workloads can compare WebGPU, including completed output readback:
make bench-jax-js-browser-model-webgpu BENCH_JAX_JS_EXTRA='--paired --rounds 4 --baseline temp/old-checkout/js/dist/polygrad.async.js'
WebGPU reports also retain shader sources, dispatch/submission counts and GPU
pass timestamps from separate diagnostic calls. Those timestamp sums exclude
host work and inter-pass gaps; they are not whole-call latency. The runner uses
the browser test launcher and fails if the requested backend cannot initialize.
Use --pipeline whole --baseline-pipeline staged to compare one whole-function
JIT with separate stage JITs in ABBA order. --pipeline captured-stages is a
WebGPU diagnostic: Polygrad captures the original materialization boundaries
in one JIT; jax-js stays staged as a control. Check saved shader/grid identities
before attributing differences to batching. --no-profile-timestamps records
host events without adding GPU queries; neither profile mode instruments the
ordinary latency samples.
Release acceptance
make test-release-list # inspect the gate list without running it
make test-release PYTHON=/path/to/test/python PARITY_PY=/path/to/parity/python \
PYTHON_MIN=/path/to/python3.9 HF_PYTHON=/path/to/hf/python \
SPECIAL_PYTHON=/path/to/scipy/python
Release acceptance runs maintained targets serially: sanitizers, backend and frontend suites, browser/Qwen, Qwen3.5/gated-delta references, independent special-function references, installed extension SDK, isolated package installs, interchange, parity, reviewed analysis, fuzzing and performance checks. CUDA and browser WebGPU are required. HIP, HCQ2/GETADDR, genuine multi-GPU and PYLITERAL are excluded, not certified. Optional MSan/TSan/Fil-C and additional browser executables have separate targets.
Required setup:
- Fetch pinned fixtures with
make fetch-llama-pretrained fetch-onnx-encoder HF_PYTHON=/path/to/hf/python, or setLLAMA_CHECKPOINTandONNX_ENCODER_DIRto existing copies. Preflight verifies their hashes; it does not download them. PYTHONandPARITY_PY: CPython 3.11 with the selected test dependencies; source-audit AST hashes require that interpreter. ReleasePYTHONdefaults toPARITY_PY.PYTHON_MIN: Python 3.9 for the isolated minimum-version package check.HF_PYTHON: the Torch/Transformers/Hugging Face and ONNX Runtime reference environment.SPECIAL_PYTHON: SciPy for independent special-function value and gradient checks.EMCCandNODE: Emscripten and Node for Wasm and installed extension SDK checks.- An AVX2/FMA CPU for the required native custom-kernel lane.
CC: defaults to Clang for release runs, with CPU-renderer__fp16support. Formatting usesCLANG_FORMAT=clang-format-14; analysis usesANALYZER_CC=clang-14at the exact reviewed version.buildinPYTHON,z3inPARITY_PY, writable caches, a browser display and the required model fixtures. OverrideQWEN3_GGUFfor its local path.BROWSER_TMPDIR: a short writable path with at least 2 GiB free for Chromium profiles and model-transfer files (default/tmp). It is separate from compiler scratch space selected byTMPDIR/POLY_TMPDIR.
Preflight stops on missing prerequisites. The runner does not download model fixtures, update baselines or approve debts. Individual package checks can use the network.
Results go to a fresh temp/release-* directory or a new RELEASE_DIR.
summary.json records gate commands, exits, times and source/log/artifact hashes;
individual test counts and skips remain in gate logs. Source changes fail the
run. Non-preflight failures do not stop later gates; any failed gate makes the
final exit nonzero. Interruptions terminate child processes and mark remaining
gates unrun. This target never publishes packages.
make test-release-runner tests orchestration without running acceptance.
make analyze is the raw zero-warning check; release acceptance uses
test-analyze-reviewed with source-bound reviews.
New warnings, stale reviews, compiler errors and incomplete scans fail.
test-symbolic-z3-supported checks general integers and division.
The separate test-symbolic-z3-fixed still fails uint8_add_wrap_cmp in both
Polygrad and pinned Tinygrad: the interval fold returns false while uint8
((250 + 10) % 256) < 5 is true. test-symbolic-z3 runs all modes.
Model and package checks
make test-release-gates # negative controls: missing HF dependencies / CUDA
make test-qwen3 # requires POLY_QWEN3_GGUF or the default temp/ fixture
make test-vision # small CLIP/ViT/DINO reference, import and native/Wasm checks
POLY_TEST_FILTER=Vision make test-browser # the same models on browser auto/INTERP/WebGPU
make HAS_CUDA=1 test-qwen3-cuda # requires a working CUDA GPU, not a skipped test
make test-hf-e2e HF_PYTHON=/path/to/python # requires huggingface_hub, transformers and torch
make test-release-packages # fresh sdist and npm native/fallback installations
HF tests can skip missing dependencies in ordinary runs; the strict target fails.
For offline fixtures, set HF_HUB_CACHE, HF_HUB_OFFLINE=1 and
TRANSFORMERS_OFFLINE=1. Package checks verify installed library origins and
native-build failure with Wasm fallback. Node browser-export checks do not
replace real-browser execution.
make fetch-llama-pretrained downloads the pinned TinyStories Llama 2-style
fixture; make test-llama-pretrained compares logits, greedy choices and bundle
reload against Transformers. Set LLAMA_TEST_DEVICES='cpu cuda' to require both.
For Llama 3.2 1B, use make test-llama32-pretrained with
LLAMA32_CHECKPOINT=/path/to/Llama-3.2-1B and
LLAMA32_TEST_DEVICES='cuda cpu'. That target does not download weights or test
KV caching/bundle export. Float32 weights alone exceed 4 GB; allow additional
host RAM/swap and device memory.
Tinygrad compatibility
The default upstream selection is backend/test_ops.py. Each output directory
must be new:
make test-upstream-runner
make test-compat-tinygrad-upstream UPSTREAM_COMPAT_DIR=temp/upstream-run-001
# Select other full files or exact pytest nodeids explicitly:
make test-compat-tinygrad-upstream UPSTREAM_COMPAT_DIR=temp/upstream-run-002 \
UPSTREAM_COMPAT_TESTS='test/null/test_dtype.py test/unit/test_conv.py'
# Reproduce the reviewed compatibility baseline (known failures remain visible):
make test-compat-tinygrad-upstream-ratchet UPSTREAM_COMPAT_DIR=temp/upstream-ratchet-001
# Separate CPU ops lane with the explicit, source-locked NIR-only adaptation:
make test-compat-tinygrad-ops UPSTREAM_COMPAT_DIR=temp/upstream-ops-001
# Ratchet the separately reviewed CPU ops baseline:
make test-compat-tinygrad-ops UPSTREAM_COMPAT_DIR=temp/upstream-ops-002 \
UPSTREAM_COMPAT_ARGS='--baseline test/fixtures/tinygrad_upstream_ops_cpu_014_baseline.json'
# NN and optimizer files, with source-locked CPU helper adaptation:
make test-compat-tinygrad-nn UPSTREAM_COMPAT_DIR=temp/upstream-nn-001
# Backend dtype diagnostics, including forced emulation (needs BF16-capable Clang):
make test-compat-tinygrad-dtype UPSTREAM_COMPAT_DIR=temp/upstream-dtype-001
# Diagnostic refresh of all selected cases per engine, in three serial lanes:
make test-compat-tinygrad-suite UPSTREAM_COMPAT_DIR=temp/upstream-suite-001
Reports record source/library hashes, provider identities, outcomes and collection errors. Missing dependencies and incomplete execution fail. CPU ops and NN lanes use source-locked helper adaptations, not replacement implementations. The CPU dtype lane extracts locked scalar test oracles and uses the core’s actual emulation controls. It preserves assertions, including the old FP8 flush-to-zero expectation superseded by approved backport PG-DIV-032; that diagnostic failure remains visible rather than weakening the test. Ratchets reject regressions and changed test/environment contracts; a passing ratchet does not mean every upstream test passes. Review baseline changes; see the runner for options.
Complete migration certification is separate: reference-migration-check
requires source-wave audits and source-bound execution evidence. Open debts are
not approvals. Durable reviews live in test/fixtures/migration/.
Performance
make test runs only the C suite. make test-browser runs deterministic Wasm
performance snapshots after its browser tests; the serial release runner lists
the snapshot as a separate gate. Node-only make test-js-wasm needs no browser.
Snapshots cover five MLP/FFN workloads with default, explicitly enabled and
explicitly disabled kernel selection.
They require Emscripten, the JS
dependencies, Playwright Chromium and Binaryen’s wasm-dis (WASM_DIS overrides
its path). Kernel hashes, decoded instruction counts, dispatches and frontend
crossings must match. Allocation counts cover exported allocation calls only;
the C suite separately checks internal LINEAR scratch allocation.
make test-wasm-performance-structure
make update-wasm-performance-structure # review the fixture diff; not a timing approval
make bench-wasm-qualify BENCH_WASM_BASELINE=evidence/baseline/polygrad.sync.js
make bench-wasm-regression BENCH_WASM_BASELINE=evidence/baseline/polygrad.sync.js
make bench-jax-js-browser-model-wasm BENCH_JAX_JS_EXTRA='--kernel-qualify --pipeline whole --kernels'
Use a separately built, frozen baseline inside the checkout (a symlink is fine).
Qualification requires same-build equivalence within 5% and detection of actual
kernel/host work added by the harness. Ten balanced fresh-page pairs per control
are fixed before running; a failed control exits 2. bench-wasm-regression
runs qualification first and stops if it fails, before comparing the candidate.
Reports retain individual samples, bundle/harness hashes and machine conditions.
Schema 2 uses the mean timed-call duration per page; compare builds in the same
run rather than comparing these numbers with older median-based reports.
Timing qualification has not passed on the shared development machine. With
ten pairs and the correction across fifteen control/workload comparisons, the
interval uses the smallest and largest ratio: one outlier can prevent a pass.
Do not use this unqualified whole-call gate to approve default-on or replace a
failed qualification with an uncorrected comparison.
--kernel-qualify checks the measuring method on one bundle, not a candidate
against a baseline. Twenty fresh pages each contain alternating identical-code
and real extra-work controls. It measures captured dispatches individually on
hot buffers and the complete dispatch sequence in order. Raw arguments remain
valid because inputs, output and JIT stay alive without intervening runtime
allocation. This mode is restricted to the maintained stateless FFN workloads.
Kernel hashes and output bits must remain unchanged. Reports keep every block,
execution counts and one aggregate ratio per page; an unresolved control exits 2.
Passing these controls does not qualify whole-call overhead or default-on.
Runtime, frontend and build changes still need timing even if kernels are identical.
These commands do not change defaults or refresh baselines automatically.
Use these owners rather than introducing another performance gate:
| Check | Role |
|---|---|
test-wasm-performance-structure |
Deterministic browser/Wasm regression check; fixture changes require review. |
bench-wasm-regression |
Frozen-build browser comparison, blocked until its timing controls qualify. |
test-release-c-performance, test-release-py-performance |
Existing paired native release checks. |
bench-smoke-regression, bench-ratios, bench-compare |
Local diagnostics; historical numbers are not browser release evidence. |
test-bench-wasm |
Tests the measurement code, not product speed. |
Browser model, stage and WebGPU profiles share the existing model benchmark.
One-off scripts under temp/ are historical evidence, not additional gates.
Use separate build checkouts for concurrent sessions: serial Make recipes do
not prevent another process from overwriting shared bundles.
make bench-local-baseline
make bench-smoke-regression
make bench-ratios
Absolute smoke baselines are machine-specific and stored under ignored paths. For paired Python eager, training and JIT/readback checks:
make bench-py-eager PY_PERF_BASELINE=/path/to/baseline/venv/bin/python
Use an idle machine and matching Python/NumPy versions. Nine alternating
candidate/baseline pairs retain all samples without retries. Every workload must
meet its 1.02 median ratio limit; numerical results and library origins are checked.
The report is temp/python-eager-performance.json.
Release acceptance creates its own isolated, hash-pinned 0.7.1 baseline install,
requiring PyPI access for dependencies. PY_PERF_BASELINE_SDIST can supply that
exact source archive locally. Evidence is under python-performance/.
The C gate separately builds the pinned 0.7.1 release tag and current source
with matching flags and nine alternating pairs. That checkpoint must exist locally.
Load averages and affinity are recorded; high load is flagged, not used to discard
samples. These guards do not certify large-model or Model API performance.
Repository
| Path | Purpose |
|---|---|
src/ |
C core, compiler, schedulers, runtimes, backends |
py/ |
Python frontend |
js/ |
Node, WASM, and browser frontend |
r/ |
Unmaintained R prototype; incompatible with the current C API and excluded from release validation |
test/ |
C tests |
License
MIT