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

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 is input or target. 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, forward exposes all declared inputs and outputs. An explicit scalar objective enables the existing Model training path. controls lists declared control names; writes maps 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;
  • always keeps the full producer graph for portable export or later placement;
  • never builds 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 set LLAMA_CHECKPOINT and ONNX_ENCODER_DIR to existing copies. Preflight verifies their hashes; it does not download them.
  • PYTHON and PARITY_PY: CPython 3.11 with the selected test dependencies; source-audit AST hashes require that interpreter. Release PYTHON defaults to PARITY_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.
  • EMCC and NODE: 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 __fp16 support. Formatting uses CLANG_FORMAT=clang-format-14; analysis uses ANALYZER_CC=clang-14 at the exact reviewed version.
  • build in PYTHON, z3 in PARITY_PY, writable caches, a browser display and the required model fixtures. Override QWEN3_GGUF for 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 by TMPDIR / 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