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