Skip to main content

Trace eager code to graphs

Eager code runs op by op. Tracing runs your function ONCE over data-free stand-ins (shapes and dtypes matter, values are never read) and captures the operator graph: a ModelGraph, the same runnable type the ONNX loader compiles to, with the same run.

Two rules make a function traceable:

  • List in, list out. A traced callable receives its input tensors as one list and returns its outputs as a list. A bare tensor return does not auto-wrap; return [y].
  • No value reads. The stand-ins carry no data, so reading a value during tracing raises. Shape-driven math is fine.
main.py
import clika_runtime as crt

w = crt.full((4, 4), 0.5)

def fn(inputs: list) -> list:
x = inputs[0]
return [(x @ w + 1.0).relu()]

x = crt.ones(4, 4)

# Capture: fn runs once over stand-ins; the graph is inspectable.
g = crt.trace(fn, example_inputs=[x])
print(g)

# The graph runs like the function did.
(y,) = g.run([x])
print(f"peak = {y.amax().item()}") # 4*0.5 + 1 = 3

Modules trace the same way: crt.trace(module, example_inputs=[x]) captures the module's forward.

The compile wrapper

compile wraps a callable in capture-and-replay: the first call runs eagerly AND captures; later calls replay the graph. A shape change recaptures, invisible to values and visible on the counter.

step = crt.compile(fn)
print(step.state) # Pending; nothing captured yet
(first,) = step([x]) # runs eagerly and captures
print(step.state) # Compiled; later calls replay
(again,) = step([x]) # replays the captured graph

step([crt.ones(2, 4)]) # a new shape recaptures transparently
print(f"recaptures: {step.recapture_count}")

step.reset() re-arms the wrapper; step.take_graph() hands over the captured ModelGraph for standalone use: save it, run it, or hand it to another part of the program.