tf.compat.v1.wrap_function

Оборачивает функцию TF 1.x fn в функцию графа.

tf.compat.v1.wrap_function(
    fn, signature, name=None
)

Используется в блокнотах

Функция Python fn будет вызвана один раз с символическими аргументами, указанными в signature, трассирована и преобразована в функцию графа. Все переменные, созданные fn, будут принадлежать объекту, возвращаемому wrap_function. Полученную функцию графа можно вызывать с тензорами, соответствующими сигнатуре.

def f(x, do_add):
  v = tf.Variable(5.0)
  if do_add:
    op = v.assign_add(x)
  else:
    op = v.assign_sub(x)
  with tf.control_dependencies([op]):
    return v.read_value()

f_add = tf.compat.v1.wrap_function(f, [tf.TensorSpec((), tf.float32), True])

assert float(f_add(1.0)) == 6.0
assert float(f_add(1.0)) == 7.0

# Can call tf.compat.v1.wrap_function again to get a new trace, a new set
# of variables, and possibly different non-template arguments.
f_sub= tf.compat.v1.wrap_function(f, [tf.TensorSpec((), tf.float32), False])

assert float(f_sub(1.0)) == 4.0
assert float(f_sub(1.0)) == 3.0

И tf.compat.v1.wrap_function, и tf.function создают вызываемый граф TensorFlow. Однако, если tf.function выполняет все операции с состоянием (например, tf.print) и упорядочивает операции, обеспечивая ту же семантику, что и при eager-выполнении, то wrap_function ближе по поведению к session.run в TensorFlow 1.x. Она не будет выполнять операции, если они не требуются для вычисления выходных данных функции — через зависимость по данным или управляющую зависимость. Операции также не будут упорядочиваться.

В отличие от tf.function, wrap_function трассирует функцию Python только один раз. Как и для плейсхолдеров в TF 1.x, для аргумента signature у wrap_function необходимо указать формы и типы данных.

Поскольку функция трассируется только один раз, переменные и состояние могут создаваться внутри функции и принадлежать объекту-обёртке функции.