Setup captures and variables in restored functions.
(self)
| 156 | coder.decode_proto(proto.canonicalized_input_signature)) |
| 157 | |
| 158 | def _setup_functions_captures(self): |
| 159 | """Setup captures and variables in restored functions.""" |
| 160 | concrete_functions = sorted(self._proto.concrete_functions.items()) |
| 161 | for name, proto in concrete_functions: |
| 162 | concrete_function = self._concrete_functions[name] |
| 163 | bound_inputs = [ |
| 164 | self._get_tensor_from_node(node_id) |
| 165 | for node_id in proto.bound_inputs] |
| 166 | bound_variables = [ |
| 167 | self._nodes[node_id] |
| 168 | for node_id in proto.bound_inputs |
| 169 | if self._proto.nodes[node_id].WhichOneof("kind") == "variable" |
| 170 | ] |
| 171 | # TODO(andresp): This is only injecting the captured inputs into the |
| 172 | # concrete function, note that we did not modify the FuncGraph |
| 173 | # itself. |
| 174 | concrete_function._captured_inputs = bound_inputs # pylint: disable=protected-access |
| 175 | concrete_function._func_graph.variables = bound_variables # pylint: disable=protected-access |
| 176 | if bound_inputs: |
| 177 | for bound_input, internal_capture in zip( |
| 178 | bound_inputs, concrete_function.inputs[-len(bound_inputs):]): |
| 179 | if ds_values.is_distributed_variable(bound_input): |
| 180 | concrete_function.graph.capture_distributed_variable( |
| 181 | bound_input, internal_capture) |
| 182 | else: |
| 183 | concrete_function.graph._captures[ops.tensor_id(bound_input)] = ( # pylint: disable=protected-access |
| 184 | bound_input, internal_capture) |
| 185 | if internal_capture.dtype == dtypes.resource: |
| 186 | if resource_variable_ops.is_resource_variable(bound_input): |
| 187 | try: |
| 188 | handle = bound_input.handle |
| 189 | except ValueError: |
| 190 | # For mirrored variables we'll copy handle data for components |
| 191 | # as they get captured. |
| 192 | pass |
| 193 | else: |
| 194 | custom_gradient.copy_handle_data(handle, internal_capture) |
| 195 | else: |
| 196 | custom_gradient.copy_handle_data(bound_input, internal_capture) |
| 197 | # Setting "captures" first means "capture" won't create a new |
| 198 | # placeholder for this input. |
| 199 | concrete_function.graph.capture(bound_input) |
| 200 | |
| 201 | def _get_tensor_from_node(self, node_id): |
| 202 | """Resolves a node id into a tensor to be captured for a function.""" |
no test coverage detected