tinygrad-notes

TinyJit: capture and replay

All tutorials · Updated September 21, 2026 · tinygrad 8ad8f73

TinyJit captures realized execution and replays it with new inputs. It is not a synonym for compiling a kernel, creating a Metal pipeline, or using CUDA graphs. Kernel compilation and runtime caching also happen without TinyJit.

Observe the three phases

from tinygrad import Tensor, TinyJit

python_calls = []
@TinyJit
def twice(x):
  python_calls.append("called")
  return (x * 2).realize()

observed = []
for value in (1.0, 2.0, 3.0, 4.0):
  result = twice(Tensor([value]).realize())
  observed.append(result.tolist())

assert observed == [[2.0], [4.0], [6.0], [8.0]]
assert len(python_calls) == 2
print(observed)
print("Python function calls:", len(python_calls))

With default JIT settings, _TinyJit.call follows this sequence:

Invocation	Python function	Execution
First	Runs	Ordinary execution and returned Tensor realization.
Second	Runs	Captures linears, combines them, lowers/plans them, and executes captured work.
Third onward	Does not run	Checks input compatibility and replays captured work.

The list append demonstrates why Python side effects must not control per-replay behavior. Logging, random choices made in Python, and Python branches do not automatically become dynamic operations in the captured graph.

What is captured?

The capture path records LINEAR graphs produced during realization. It combines their calls into a larger linear graph and invokes jit_lower. The implementation accounts for held buffers before memory planning so that externally live storage is not incorrectly reused.

CapturedJit retains the return structure and executable graph. Each replay supplies current input buffer UOps and symbolic bindings. Where supported, backend graph execution can batch launches. The benefit is reducing repeated Python scheduling/launch overhead; it is distinct from making one kernel’s arithmetic faster.

Inputs must fit the captured computation

The replay path compares input names and expected input information. This is not an arbitrary-shape tracing cache that silently builds a new graph for every Python invocation. Changing shape, layout, device, or other captured input requirements can invalidate replay.

Realize new input tensors before calling the JIT in small experiments, as above. That makes the input storage explicit and keeps construction outside the capture being studied. Avoid reading results with item() or tolist() inside the captured function; keep host inspection outside.

Nested active TinyJit capture is explicitly rejected. Calling another captured function and building a nested capture are different situations; consult the implementation before assuming nesting is supported.

Output storage and timing

Check replay values and reject a changed shape

from tinygrad import Tensor, TinyJit
from tinygrad.engine.jit import JitError

@TinyJit
def rows(x):
  return (x.T.pad(((1, 1), (1, 1))) + 10).sum(axis=1).realize()

snapshots = []
for offset in (0., 1., -2., 4.):
  data = [[1.+offset, 2.+offset, 3.+offset], [4.+offset, 5.+offset, 6.+offset]]
  snapshots.append(rows(Tensor(data).realize()).tolist())
  assert snapshots[-1] == [40., 45.+2*offset, 47.+2*offset, 49.+2*offset, 40.]
assert snapshots[0] == [40., 45., 47., 49., 40.]
try:
  rows(Tensor([[1., 2.], [3., 4.]]).realize())
except JitError as error:
  assert "args mismatch" in str(error)
else:
  raise AssertionError("shape-changing replay was unexpectedly accepted")

This checks two replay invocations with changed data and an incompatible shape after capture. The saved Python lists remain snapshots; retaining returned Tensor objects would not provide the same guarantee. Rejection here is expected behavior, not a request to silently retrace.

Captured outputs can reuse storage across replays. If you need a historical result, copy its values before the next invocation; the example stores Python lists immediately. Saving Tensor references alone does not establish independent snapshots.

For timing, warm up through capture, synchronize the device before and after measured work, and distinguish end-to-end latency from kernel timings. Compilation, capture, and replay are separate costs.

The current reset() method clears capture state. Resetting is an explicit request to build a new captured execution, not evidence that arbitrary varying inputs were valid for the old one.

Read profiling for a synchronized timing example and command queues for device submission.

Original chapter by Di Zhu: historical version.