tf.types.experimental.TraceType

Представляет тип объекта (объектов) для трассировки tf.function.

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

TraceType — это абстрактный класс, от которого могут наследоваться другие классы, чтобы предоставлять сведения о связанных классах для целей трассировки tf.function. Логика типизации, реализуемая этим механизмом, используется для принятия решений об использовании кэшированных конкретных функций и повторной трассировке.

Например, рассмотрим следующую tf.function и классы:

@tf.function
def get_mixed_flavor(fruit_a, fruit_b):
  return fruit_a.flavor + fruit_b.flavor

class Fruit:
  flavor = tf.constant([0, 0])

class Apple(Fruit):
  flavor = tf.constant([1, 2])

class Mango(Fruit):
  flavor = tf.constant([3, 4])

tf.function не знает, когда следует повторно использовать существующую конкретную функцию для класса Fruit, поэтому при простом подходе она выполняет трассировку заново для каждого нового экземпляра.

get_mixed_flavor(Apple(), Mango()) # Traces a new concrete function
get_mixed_flavor(Apple(), Mango()) # Traces a new concrete function again

Однако мы, разработчики класса Fruit, знаем, что у каждого подкласса есть фиксированный вариант, и можем повторно использовать существующую трассированную конкретную функцию, если подкласс совпадает. Отказ от такой избыточной трассировки конкретных функций может значительно повысить производительность.

class FruitTraceType(tf.types.experimental.TraceType):
  def __init__(self, fruit):
    self.fruit_type = type(fruit)
    self.fruit_value = fruit

  def is_subtype_of(self, other):
     return (type(other) is FruitTraceType and
             self.fruit_type is other.fruit_type)

  def most_specific_common_supertype(self, others):
     return self if all(self == other for other in others) else None

  def placeholder_value(self, placeholder_context=None):
    return self.fruit_value

class Fruit:

 def __tf_tracing_type__(self, context):
   return FruitTraceType(self)

Теперь попробуем вызвать её ещё раз:

get_mixed_flavor(Apple(), Mango()) # Traces a new concrete function
get_mixed_flavor(Apple(), Mango()) # Re-uses the traced concrete function

Методы

cast

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

cast(
    value, cast_context
) -> Any

Приводит значение к этому типу.

flatten

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

flatten() -> List['TraceType']

Возвращает список TensorSpec, соответствующих значениям to_tensors.

from_tensors

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

from_tensors(
    tensors: Iterator[core.Tensor]
) -> Any

Создаёт значение этого типа из тензоров.

Количество тензоров должно быть таким же фиксированным, как и для to_tensors.

is_subtype_of

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

@abc.abstractmethod
is_subtype_of(
    other: 'TraceType'
) -> bool

Возвращает True, если self является подтипом other.

Например, tf.function использует подтипизацию для диспетчеризации: если a.is_subtype_of(b) равно True, аргумент типа TraceType a можно использовать как аргумент для ConcreteFunction, трассированной с TraceType b.

Пример:

class Dimension(TraceType):
  def __init__(self, value: Optional[int]):
    self.value = value

  def is_subtype_of(self, other):
    # Either the value is the same or other has a generalized value that
    # can represent any specific ones.
    return (self.value == other.value) or (other.value is None)

most_specific_common_supertype

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

@abc.abstractmethod
most_specific_common_supertype(
    others: Sequence['TraceType']
) -> Optional['TraceType']

Возвращает наиболее конкретный супертип self и others, если он существует.

Возвращаемый TraceType является супертипом self и others, то есть все они являются его подтипами (см. is_subtype_of). Кроме того, он наиболее конкретен: у него нет подтипа, который также являлся бы общим супертипом self и others.

Если у self и others нет общего супертипа, возвращается None.

Пример:

class Dimension(TraceType):
  def __init__(self, value: Optional[int]):
    self.value = value

  def most_specific_common_supertype(self, other):
     # Either the value is the same or other has a generalized value that
     # can represent any specific ones.
     if self.value == other.value:
       return self.value
     else:
       return Dimension(None)

placeholder_value

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

@abc.abstractmethod
placeholder_value(
    placeholder_context
) -> Any

Создаёт заглушку для трассировки.

tf.funcion выполняет трассировку с использованием значения-заглушки, а не фактического значения. Например, значение-заглушка может представлять несколько разных фактических значений. Это означает, что трассировка, созданная с таким значением-заглушкой, будет более обобщённой и пригодной для повторного использования, что позволит избежать затратной повторной трассировки.

Для приведённого выше примера с Fruit реализация:

class FruitTraceType:
  def placeholder_value(self, placeholder_context):
    return Fruit()

указывает tf.function выполнять трассировку с объектами Fruit() вместо фактических объектов Apple() и Mango() при вызове get_mixed_flavor(Apple(), Mango()). Например, аргументы Tensor заменяются тензорами с похожими формой и типом данных, полученными на выходе операции tf.Placeholder.

В более общем смысле значения-заглушки — это аргументы tf.function такими, какими они видны из тела функции:

@tf.function
def foo(x):
  # Here `x` is be the placeholder value
  ...

foo(x) # Here `x` is the actual value

to_tensors

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

to_tensors(
    value: Any
) -> List[core.Tensor]

Разбивает значение этого типа на тензоры.

Для экземпляра TraceType количество тензоров, создаваемых для соответствующего значения, должно быть постоянным.

__eq__

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

@abc.abstractmethod
__eq__(
    other
) -> bool

Возвращает результат сравнения self==value.