Merge a node def without materializing a full DeviceSpec object. Often a device merge is invoked in order to generate a string which can be passed into the c api. In such a case, we can cache the node_def.device -> merge_result_string map, and in most cases avoid: - Mater
(self, node_def)
| 131 | return self._spec.make_merged_spec(current_device) |
| 132 | |
| 133 | def shortcut_string_merge(self, node_def): |
| 134 | """Merge a node def without materializing a full DeviceSpec object. |
| 135 | |
| 136 | Often a device merge is invoked in order to generate a string which can be |
| 137 | passed into the c api. In such a case, we can cache the |
| 138 | node_def.device -> merge_result_string |
| 139 | |
| 140 | map, and in most cases avoid: |
| 141 | - Materializing a copy of self._spec (In the case of DeviceSpecV1) |
| 142 | - Materializing a DeviceSpec for node_def.device |
| 143 | - A DeviceSpec.merge_from invocation |
| 144 | |
| 145 | In practice the cache hit rate for this function is very high, because the |
| 146 | number of invocations when iterating through the device stack is much |
| 147 | larger than the number of devices. |
| 148 | |
| 149 | Args: |
| 150 | node_def: An Operation (or Operation-like) to merge device constraints |
| 151 | with self._spec |
| 152 | |
| 153 | Returns: |
| 154 | A string containing the merged device specification. |
| 155 | """ |
| 156 | device = node_def.device or "" |
| 157 | |
| 158 | merge_key = (self._spec, device) |
| 159 | result = _string_merge_cache.get(merge_key) |
| 160 | if result is None: |
| 161 | # This update is not atomic, however because the merge is stateless |
| 162 | # we don't need to lock when updating the cache. |
| 163 | result = self.__call__(node_def).to_string() |
| 164 | _string_merge_cache[merge_key] = result |
| 165 | |
| 166 | return result |
| 167 | |
| 168 | def __repr__(self): |
| 169 | return "{} (spec: {})".format( |
no test coverage detected