tf.wrap_function

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

tf.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) и операции последовательности, чтобы обеспечить такие же семантики, как при нетерпеливом выполнении, wrap_function ближе к поведению session.run в TensorFlow 1.x. Он не будет выполнять какие-либо операции, пока они не потребуются для вычисления выходных данных функции, либо через зависимость данных, либо через зависимость управления. Он также не будет упорядочивать операции.

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

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