Что такое @tf.function и каковы его ограничения (побочные эффекты Python, трейсинг)?

Механизм трейсинга @tf.function

При первом вызове декорированной функции TF запускает трейсинг: исполняет Python-тело, записывает TF-операции в граф и кеширует результат как ConcreteFunction. Последующие вызовы с теми же сигнатурами аргументов используют кеш, минуя Python-интерпретатор. Это источник как производительности, так и большинства ловушек.

Ограничение 1: Python side-effects работают только при трейсинге

import tensorflow as tf

call_count = 0

@tf.function
def f(x):
    global call_count
    call_count += 1          # Python side-effect!
    print(f"Tracing #{call_count}")  # тоже только при трейсинге
    return x * 2

f(tf.constant(1))  # Tracing #1 — трейсинг
f(tf.constant(2))  # Молчит — граф уже закеширован
print(call_count)  # 1, а не 2!

# Правильно: используйте tf.print для логирования внутри графа
@tf.function
def f_correct(x):
    tf.print("Value:", x)  # выполняется при каждом вызове графа
    return x * 2

Ограничение 2: повторный трейсинг при изменении формы/типа

@tf.function
def compute(x):
    print(f"Tracing: dtype={x.dtype}, shape={x.shape}")
    return tf.reduce_sum(x)

compute(tf.constant([1.0, 2.0]))       # Tracing: shape=(2,)
compute(tf.constant([1.0, 2.0, 3.0])) # Tracing: shape=(3,) — новый граф!
compute(tf.constant(1.0))             # Tracing: shape=() — третий граф!

# Решение: зафиксировать сигнатуру
@tf.function(input_signature=[tf.TensorSpec([None], tf.float32)])
def compute_fixed(x):
    return tf.reduce_sum(x)

# Теперь любой 1D float32 тензор — один граф
compute_fixed(tf.constant([1.0, 2.0]))
compute_fixed(tf.constant([1.0, 2.0, 3.0]))  # без трейсинга

Ограничение 3: Python-скаляры как аргументы вызывают трейсинг для каждого значения

@tf.function
def scale(x, factor):  # factor — Python int
    return x * factor

scale(tf.constant([1.0]), 2)  # трейсинг для factor=2
scale(tf.constant([1.0]), 3)  # трейсинг для factor=3, НОВЫЙ граф!
scale(tf.constant([1.0]), 4)  # ещё один...

# Решение: передавайте Python-константы как tf.Tensor
@tf.function
def scale_tensor(x, factor):  # factor — тензор
    return x * factor

scale_tensor(tf.constant([1.0]), tf.constant(2))
scale_tensor(tf.constant([1.0]), tf.constant(3))  # без трейсинга

Ограничение 4: произвольные Python-объекты не сериализуются в граф

import numpy as np

@tf.function
def bad_fn(x):
    arr = np.array([1, 2, 3])  # NumPy создаётся при трейсинге и "вмораживается"
    return x + tf.constant(arr)  # OK, но arr фиксируется на момент трейсинга

# Изменение arr снаружи НЕ повлияет на граф после первого трейсинга
# Решение: передавайте как аргумент или конвертируйте явно

Ограничение 5: циклы и условия с Python-значениями

@tf.function
def loop_python(x, n):  # n — Python int
    for i in range(n):  # разворачивается в n отдельных операций при трейсинге!
        x = x + 1
    return x

# При n=1000 граф содержит 1000 операций сложения
# Решение: используйте tf.while_loop для динамических итераций
@tf.function
def loop_tf(x, n):  # n — tf.Tensor
    return tf.while_loop(
        cond=lambda i, _: i < n,
        body=lambda i, v: (i + 1, v + 1),
        loop_vars=[tf.constant(0), x],
    )[1]

Подводные камни

• Утечка памяти при частом трейсинге: каждый ConcreteFunction хранит граф в памяти. При динамическом создании функций (tf.function внутри цикла) графы накапливаются.
• Использование tf.Variable, созданной внутри @tf.function, при повторном вызове вызывает ошибку — переменные должны создаваться вне функции.
• Исключения Python внутри @tf.function превращаются в ошибки TF и теряют Python-трейсбэк — для отладки временно убирайте декоратор.
• Предупреждение «5 out of the last 5 calls to <fn> triggered tf.function retracing» — сигнал о чрезмерном трейсинге; проверяйте типы и формы аргументов.
• Вложенные @tf.function трейсятся независимо — outer-функция не «видит» inner-граф напрямую.
• Генераторы и итераторы Python нельзя передать в граф — используйте tf.data.Dataset.
• При использовании jit_compile=True (XLA) динамические формы не поддерживаются — все операции должны иметь статические или частично статические размеры.

• Объяснять tf function limitations только синтаксисом без shape, dtype, состояния или режима выполнения.
• Игнорировать leakage, воспроизводимость, пустые входы и скрытые копии данных.
• Не проверять production-симптомы: latency, память, ретраи, дрейф качества и несовпадение версий.

• Может ли связать tf function limitations с реальным контрактом входов и выходов.
• Упоминает ли тесты, метрики, reproducibility и диагностику ошибок.
• Видит ли различие между demo-кодом в ноутбуке и production-пайплайном.
• Предлагает ли observability, rollback, ограничения стоимости и стратегию incident replay.