tf.types.experimental.TraceType

Represents the type of object(s) for tf.function tracing purposes.

TraceType is an abstract class that other classes might inherit from to provide information regarding associated class(es) for the purposes of tf.function tracing. The typing logic provided through this mechanism will be used to make decisions regarding usage of cached concrete functions and retracing.

For example, if we have the following tf.function and classes:

@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 does not know when to re-use an existing concrete function in regards to the Fruit class so naively it retraces for every new instance.

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

However, we, as the designers of the Fruit class, know that each subclass has a fixed flavor and we can reuse an existing traced concrete function if it was the same subclass. Avoiding such unnecessary tracing of concrete functions can have significant performance benefits.

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

  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

class Fruit:

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

Now if we try calling it again:

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

Methods

is_subtype_of

View source

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

Returns True if self is a subtype of other.

For example, tf.function uses subtyping for dispatch: if a.is_subtype_of(b) is True, then an argument of TraceType a can be used as argument to a ConcreteFunction traced with an a TraceType b.

Args
other A TraceType object to be compared against.

Example:

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

View source

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

Returns the most specific supertype of self and others, if exists.

The returned TraceType is a supertype of self and others, that is, they are all subtypes (see is_subtype_of) of it. It is also most specific, that is, there it has no subtype that is also a common supertype of self and others.

If self and others have no common supertype, this returns None.

Args
others A sequence of TraceTypes.

Example:

 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)

__eq__

View source

@abc.abstractmethod
__eq__(
    other
) -> bool

Return self==value.