Lab 03 — TF tf.function Tracing & Graph Analysis

Phase: 01 — PyTorch & TF Framework Internals | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–4 hours

tf.function is the boundary between eager Python and TensorFlow's graph runtime. Every "the model behaves differently in graph mode" bug lives at this boundary. This lab makes the tracing machinery visible.

What you build

• trace_and_inspect — capture the ConcreteFunction for given inputs and produce a GraphReport: node counts by op type, captured external tensors, structured input signature
• RetraceProbe — a wrapper that counts how many times Python tracing actually runs, used to demonstrate (and then prevent) the classic retracing triggers
• stabilize — apply an explicit input_signature so shape/value changes reuse one trace
• autograph_diff — show the AutoGraph-transformed source of a Python function with if/for control flow (what your Python actually becomes)
• roundtrip_savedmodel — save → reload → verify outputs match exactly
• convert_tflite_dynamic — TFLite dynamic-range quantization with max-abs output delta vs the TF model on random probe inputs

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v          # validate your lab.py
python solution.py             # run the full demo: retrace report + TFLite delta

Suggested TODO order

1. trace_and_inspect — everything else builds on being able to see a FuncGraph
2. RetraceProbe — prove to yourself exactly which calls retrace and why
3. stabilize — make the retraces stop with an input_signature
4. autograph_diff — look at what AutoGraph did to your if statement
5. roundtrip_savedmodel, then convert_tflite_dynamic

Success criteria

• You can enumerate the retracing triggers from memory and predict the retrace count of a call sequence before running it
• You can explain why a print() inside tf.function fires once but tf.print fires every call
• Your TFLite dynamic-range conversion shows max output delta < 1e-2 on the small MLP (weights-only quantization on a well-conditioned model is nearly lossless)

Interview Q&A

Q: When does tf.function retrace? New input dtype; new rank or (without a signature) new concrete shape; a new Python-native argument value (every distinct Python int/str is a new signature!); a new closure/object identity for weakly-referenced captures; calling with kwargs vs args inconsistently.

Q: Why did my counter only increment once even though I call the function in a loop? The increment is a Python side effect — it runs at trace time. The traced graph contains only TF ops. Use tf.Variable.assign_add to make mutation part of the graph.

Q: Why is passing a Python int as batch_size a performance bug? Each distinct int is a different trace signature → unbounded retracing. Pass it as a tf.Tensor or pin an input_signature with tf.TensorSpec([None, ...]).

Q: When is TFLite conversion not lossless? Dynamic-range quantization perturbs weights (INT8); ops without TFLite kernels get rewritten (e.g., erf-based GELU → tanh approximation); fused kernels may use different accumulation order; resource variables get frozen to constants.

Extensions

• Repeat the analysis on BERT-tiny and evaluate the MNLI accuracy delta after dynamic-range quantization (the full-scale version of this lab, see phase README)
• Inspect the XLA HLO with tf.function(jit_compile=True) and fn.experimental_get_compiler_ir(...)('hlo')
• Profile trace time vs execution time with the TF Profiler

References

• tf.function guide — the retracing rules, authoritative
• Introduction to graphs and tf.function
• AutoGraph reference
• TFLite post-training quantization