def is_subtype_of(self, other: trace.TraceType) -> bool:
"""Returns True if `self` is a subtype of `other`.
Implements the tf.types.experimental.func.TraceType interface.
If not overridden by a subclass, the default behavior is to assume the
TypeSpec is covariant upon attributes that implement TraceType and
invariant upon rest of the attributes as well as the structure and type
of the TypeSpec.
Args:
other: A TraceType object.
"""
if type(self) is not type(other):
return False
is_subtype = True
def check_attribute(attribute_self, attribute_other):
nonlocal is_subtype
if not is_subtype:
return
if isinstance(attribute_self, trace.TraceType):
if not attribute_self.is_subtype_of(attribute_other):
is_subtype = False
return
else:
if attribute_self != attribute_other:
is_subtype = False
try:
# TODO(b/217959193): Replace _serialize with parameter decomposition.
nest.map_structure(check_attribute, self._serialize(),
other._serialize()) # pylint: disable=protected-access
except (ValueError, TypeError):
return False
return is_subtype
def dispatch(self, request: trace.TraceType) -> Optional[trace.TraceType]:
"""Returns the deepest subtype target if it exists in the table."""
# For known exact matches.
if request in self._dispatch_table:
return request
# For known non-exact matches.
# (self._dispatch cache does not contain exact matches)
if request in self._dispatch_cache:
# Move to the front of LRU cache.
result = self._dispatch_cache.pop(request)
self._dispatch_cache[request] = result
return result
most_specific_subtype = None
for other in self._dispatch_table:
if request.is_subtype_of(other):
if most_specific_subtype is None or other.is_subtype_of(
most_specific_subtype):
most_specific_subtype = other
self._cache_dispatch(request, most_specific_subtype)
return most_specific_subtype
def add_target(self, target: trace.TraceType) -> None:
"""Adds a new target type."""
self._dispatch_table[target] = None
for request in self._dispatch_cache:
if target.is_subtype_of(self._dispatch_cache[request]):
self._dispatch_cache[request] = target