tf.types.experimental.GenericFunction

Базовый класс для полиморфных функций графа.

Наследуется от: Callable

Функции графа — это объекты Python, вызываемые по умолчанию, которые передают вызовы в граф TensorFlow. Полиморфные функции графа могут быть основаны на нескольких графах TF и автоматически выбирать соответствующую специализацию на основе типа входных данных, с которыми они были вызваны. Они также могут создавать специализации на лету, если это необходимо, например, посредством трассировки.

См. также tf.function.

Методы

experimental_get_compiler_ir

Просмотреть исходный код

experimental_get_compiler_ir(
    *args, **kwargs
)

Возвращает IR компилятора для скомпилированной функции.

Этот API предназначен только для отладки, так как нет гарантий обратной совместимости возвращаемого IR или допустимых значений stage.

@tf.function(jit_compile=True)
def f(x):
  return x + 1

f.experimental_get_compiler_ir(tf.random.normal([10, 10])(stage='hlo')

HloModule a_inference_f_13__.9

ENTRY %a_inference_f_13__.9 (arg0.1: f32[10,10]) -> f32[10,10] {
  %arg0.1 = f32[10,10]{1,0} parameter(0), parameter_replication={false}
  %reshape.2 = f32[10,10]{1,0} reshape(f32[10,10]{1,0} %arg0.1)
  %constant.3 = f32[] constant(1)
  %broadcast.4 = f32[10,10]{1,0} broadcast(f32[] %constant.3)
  %add.5 = f32[10,10]{1,0} add(f32[10,10]{1,0} %reshape.2,
                               f32[10,10]{1,0} %broadcast.4)
  %reshape.6 = f32[10,10]{1,0} reshape(f32[10,10]{1,0} %add.5)
  %tuple.7 = (f32[10,10]{1,0}) tuple(f32[10,10]{1,0} %reshape.6)
  ROOT %get-tuple-element.8 = f32[10,10]{1,0}
    get-tuple-element((f32[10,10]{1,0}) %tuple.7), index=0
}

get_concrete_function

Просмотреть исходный код

get_concrete_function(
    *args, **kwargs
) -> tf.types.experimental.ConcreteFunction

Возвращает ConcreteFunction специализированную по типам входных данных.

Аргументы, указанные args и kwargs, следуют обычным правилам вызова функции. Возвращённая ConcreteFunction имеет тот же набор позиционных и ключевых аргументов, что и self, но их типы совместимы с типами, указанными в args и kwargs (хотя и не обязательно равны).

@tf.function
def f(x):
  return x
f_concrete = f.get_concrete_function(tf.constant(1.0))
f_concrete = f.get_concrete_function(x=tf.constant(1.0))

В отличие от обычных вызовов, get_concrete_function допускают указание типов вместо объектов TensorFlow, так, например, tf.Tensor могут быть заменены на tf.TensorSpec.

@tf.function
def f(x):
  return x
f_concrete = f.get_concrete_function(tf.TensorSpec([], tf.float64))

Если определение функции допускает только одну специализацию, args и kwargs могут быть опущены полностью.

@tf.function(input_signature=[tf.TensorSpec(None, tf.float32)])
def f(x):
  return x
f_concrete = f.get_concrete_function()

Возвращённая ConcreteFunction может вызываться обычным способом:

f_concrete(tf.constant(1.0))
<tf.Tensor: shape=(), dtype=float32, numpy=1.0>
f_concrete(x=tf.constant(1.0))
<tf.Tensor: shape=(), dtype=float32, numpy=1.0>

__call__

Просмотреть исходный код

__call__(
    *args, **kwargs
)

Выполняет этот вызов.

Это работает как обычная операция — в режиме eager она немедленно запускает выполнение, возвращая результаты. В режиме графа она создает операции, которые возвращают символические значения TensorFlow (такие как tf.Tensor, tf.data.Dataset и т. д.). Например, вызываемые объекты tf.function обычно генерируют операцию tf.raw_ops.PartitionedCall, но не всегда — точные генерируемые операции являются внутренней реализацией.