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, формы и типы данных должны быть предоставлены аргументу wrap_function в signature.

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