tf.function and AutoGraph

The mechanism that lets users write Pythonic code with regular if/while/for, decorate it with @tf.function, and get a graph-mode speedup. Spread across tensorflow/python/eager/polymorphic_function/ (the tracing engine), tensorflow/python/autograph/ (the source rewriter), and the runtime function library.

Purpose

• Convert a Python function into a FuncGraph (a tf.Graph produced by tracing).
• Cache the traced graph per input signature so repeated calls reuse it.
• Translate Python control flow (if, while, for, break, continue) into TF control-flow ops (tf.cond, tf.while_loop).
• Hook variables, captures, and side effects into the resulting graph.

Directory layout

tensorflow/python/eager/polymorphic_function/
├── polymorphic_function.py          # @tf.function user API
├── concrete_function.py             # ConcreteFunction
├── tracing_compilation.py           # First-call tracing
├── monomorphic_function.py
├── function_type_utils.py
├── attributes.py
├── compiler_ir.py                   # for jit_compile=True
└── ...

tensorflow/python/autograph/
├── core/                            # Conversion driver, error remapping
├── converters/                      # One file per Python construct
│   ├── conditional_expressions.py
│   ├── continue_statements.py
│   ├── control_flow.py
│   ├── functions.py
│   ├── lists.py
│   ├── return_statements.py
│   └── ...
├── pyct/                            # Python AST manipulation utilities
├── operators/                       # Runtime helpers (e.g. py_for_stmt, py_if_stmt)
├── lang/                            # AutoGraph-specific syntax (`tf.autograph.experimental.do_not_convert`, etc.)
├── impl/                            # Top-level conversion API
└── utils/

How a @tf.function call is dispatched

sequenceDiagram
    participant User
    participant tf.function as Function
    participant Cache as ConcreteFunction cache
    participant AG as AutoGraph
    participant Trace as Tracer
    participant FuncLib as EagerContext function library

    User->>Function: f(x)
    Function->>Cache: lookup signature(x.shape, x.dtype, ...)
    alt cached
        Cache-->>Function: cached ConcreteFunction
    else first call
        Function->>AG: optionally rewrite Python source
        AG-->>Function: rewritten code
        Function->>Trace: run rewritten code with symbolic Tensors
        Trace->>FuncLib: register FuncGraph as a function
        Trace-->>Function: ConcreteFunction
        Function->>Cache: insert(signature, ConcreteFunction)
    end
    Function->>FuncLib: execute ConcreteFunction(real x)
    FuncLib-->>User: outputs

AutoGraph

AutoGraph (tensorflow/python/autograph/) is a Python-source-to-Python-source converter. It reads the function's source via inspect.getsource, parses it with the gast library (vendored under pyct/), and rewrites Python constructs into TF-friendly equivalents:

• if x > 0: ... else: ... → tf.cond(x > 0, lambda: ..., lambda: ...) when x is a tf.Tensor.
• while predicate(...): body(...) → tf.while_loop.
• for x in dataset: body(x) → a tf.data.Dataset.reduce or tf.while_loop over the iterator.
• break/continue are translated to predicate adjustments.

When the loop variables are Python (not TF), AutoGraph leaves them alone. The decision is made by tensorflow/python/autograph/converters/control_flow.py based on the runtime types.

AutoGraph is on by default for @tf.function; pass autograph=False to disable.

Tracing

The traced execution is the standard graph-mode build path: every TF op call goes into the current FuncGraph instead of running. Captures (Python variables outside the function, tf.Variables) become inputs to the function. The result is a FuncGraph that gets registered with the EagerContext function library and called like an op.

ConcreteFunction (tensorflow/python/eager/polymorphic_function/concrete_function.py) is the wrapper around a registered FuncGraph plus its input/output signatures.

Polymorphism

@tf.function is polymorphic: each new input signature triggers a new trace. The signature includes shapes, dtypes, and Python-value identities of non-tensor arguments. The cache lives per-Function instance.

tf.function(input_signature=…) pins a single signature, preventing retracing. tf.function(reduce_retracing=True) heuristically merges similar signatures.

jit_compile=True

@tf.function(jit_compile=True) forces the traced FuncGraph to be lowered to HLO (via tf2xla) and compiled by XLA. compiler_ir.py exposes the resulting HLO and post-optimisation IR for inspection (used by tf.function.experimental_get_compiler_ir(...)).

Variable handling

A tf.Variable accessed inside a @tf.function is captured as a function input the first call. Variable creation inside the function is a common pitfall — TF will raise on the second call. The recommended pattern is tf.init_scope() or hoisting variable creation to the enclosing class' __init__.

Side effects

Side effects (assignments, queue ops, tf.print) are tracked via control dependencies inserted into the traced graph. The @tf.function decorator preserves "Python-order" semantics for these by chaining them on a hidden control edge.

Integration points

• Eager runtime (tensorflow/core/common_runtime/eager/) — registers and executes ConcreteFunctions.
• Function library in tensorflow/core/framework/function.h — where FuncGraphs live.
• Grappler — runs over the traced FuncGraph before execution.
• XLA — jit_compile=True routes to tf2xla.

Entry points for modification

• AutoGraph rewrites — tensorflow/python/autograph/converters/.
• Tracing semantics — tensorflow/python/eager/polymorphic_function/.
• Variable capture — polymorphic_function/concrete_function.py and tensorflow/python/ops/resource_variable_ops.py.

Related

• systems/eager-execution — the runtime layer that hosts the function library.
• compilers/xla — jit_compile=True.