Named tuple used to key the function cache.
| 77 | |
| 78 | |
| 79 | class CacheKey( |
| 80 | collections.namedtuple("CacheKey", [ |
| 81 | "input_signature", "parent_graph", "device_functions", |
| 82 | "colocation_stack", "in_cross_replica_context" |
| 83 | ])): |
| 84 | """Named tuple used to key the function cache.""" |
| 85 | |
| 86 | def __hash__(self): |
| 87 | """Provide a hash even if the input signature objects aren't hashable.""" |
| 88 | return hash(self._fields_safe) |
| 89 | |
| 90 | @property |
| 91 | def _fields_safe(self): |
| 92 | """Hash & equality-safe version of all the namedtuple fields.""" |
| 93 | return (self._hash_fix(self.input_signature), self.parent_graph, |
| 94 | self.device_functions, self.colocation_stack, |
| 95 | self.in_cross_replica_context) |
| 96 | |
| 97 | def _hash_fix(self, elem): |
| 98 | """Ensure elem is hashable even if a Variable is nested in it.""" |
| 99 | # Descend into tuples |
| 100 | if isinstance(elem, tuple): |
| 101 | return tuple(self._hash_fix(i) for i in elem) |
| 102 | |
| 103 | if isinstance(elem, set): |
| 104 | return {self._hash_fix(i) for i in elem} |
| 105 | |
| 106 | # If the element is not hashable, assume it is a weakref to a variable and |
| 107 | # return the dtype & shape. Else, simply return the element |
| 108 | try: |
| 109 | hash(elem) |
| 110 | except TypeError: |
| 111 | v = elem() |
| 112 | return (v.__class__, tensor_spec.TensorSpec(v.shape, v.dtype)) |
| 113 | |
| 114 | return elem |
| 115 | |
| 116 | def __eq__(self, other): |
| 117 | return self._fields_safe == other._fields_safe # pylint: disable=protected-access |
| 118 | |
| 119 | |
| 120 | CacheKey.replace = CacheKey._replace # pylint: disable=protected-access |