Skip to main content

CompiledFunction

A callable whose first call serves the eager result and records the graph; later calls with the same input structure and matching shapes run the graph without calling Python.

Made by :func:`compile`. The call takes and returns pytrees: tensor
leaves feed the graph, every other leaf is a static value baked into
the capture, so a call whose structure or static values differ from the
capture re-captures (one capture per structure is kept, the oldest
dropped past a small count). Shapes are the graph's own guard: a call
whose shapes drift from a static capture serves eagerly and re-captures
with the drifted axes dynamic; past a small budget of re-captures the
function serves eagerly until :meth:`reset`.

``state`` reports the current capture (``State.Pending``,
``State.Compiled`` or ``State.Fallback``), ``recapture_count`` how many
times a guard miss or a shape drift forced a new capture, and
:meth:`take_graph` hands out the captured graph.

input_names (property)​

The graph input names of the current capture, in flattening order.

recapture_count (property)​

How many times a changed input structure, a changed static value, or drifted shapes forced a new capture.

state (property)​

Where the current capture is in its lifecycle: State.Pending (nothing captured yet), State.Compiled (the graph serves matching calls) or State.Fallback (serving eagerly until :meth:reset).

State​

State(value, names=None, *, module=None, qualname=None, type=None, start=1)

Where the function is in its capture lifecycle: Pending (no graph yet), Compiled (a captured graph serves shape-matching calls), Fallback (the function serves eagerly and will not re-capture on its own).

__init__​

__init__(self, target: 'Callable[P, R]', *, signature: 'Sequence[TensorSpec]' = (), dynamic: 'bool | None' = None, fullgraph: 'bool' = False, optimize: 'bool | None' = None, specialize: 'bool | None' = None, place: 'PlaceLike' = None, dynamic_axes: 'DynamicAxes | None' = None, output_names: 'Sequence[str] | None' = None, graph_name: 'str | None' = None) -> 'None'

Initialize self. See help(type(self)) for accurate signature.

reset​

reset(self) -> 'None'

reset() -> None

Drop every capture and the re-capture count; the next call captures again.

take_graph​

take_graph(self) -> 'object'

take_graph() -> ModelGraph

Take the captured graph out of the current capture (State.Compiled only); the function then serves eagerly until :meth:reset.