Internal wrap function method with extended func_graph arguments.
(self,
fn,
args=None,
kwargs=None,
signature=None,
name=None)
| 500 | return self._wrap_function(fn, signature=signature, name=name) |
| 501 | |
| 502 | def _wrap_function(self, |
| 503 | fn, |
| 504 | args=None, |
| 505 | kwargs=None, |
| 506 | signature=None, |
| 507 | name=None): |
| 508 | """Internal wrap function method with extended func_graph arguments.""" |
| 509 | fn_with_filter_and_scope, returned_ops = _filter_returned_ops( |
| 510 | self._variable_holder.call_with_variable_creator_scope(fn)) |
| 511 | |
| 512 | func_graph.func_graph_from_py_func( |
| 513 | None, # Name is unused. |
| 514 | fn_with_filter_and_scope, |
| 515 | args=args, |
| 516 | kwargs=kwargs, |
| 517 | signature=signature, |
| 518 | add_control_dependencies=False, |
| 519 | func_graph=self.graph) |
| 520 | |
| 521 | # This code relies on questional behavior from `func_graph_from_py_func`. |
| 522 | # If an existing FuncGraph is passed into the `func_graph` arg, the inputs |
| 523 | # and structured outputs are overwritten. Pretty sure this is a bug, |
| 524 | # because structured outputs doesn't match up with the outputs... |
| 525 | fn_inputs = self.graph.inputs[:-len(self.graph.captures)] |
| 526 | |
| 527 | # Return filtered ops to the flattened outputs. |
| 528 | flat_fn_outputs = nest.flatten(self.graph.structured_outputs) |
| 529 | for index, op in returned_ops.items(): |
| 530 | flat_fn_outputs[index] = op |
| 531 | fn_outputs = nest.pack_sequence_as(self.graph.structured_outputs, |
| 532 | flat_fn_outputs) |
| 533 | |
| 534 | name = name or fn.__name__ |
| 535 | wrapped_function = self._wrapped_function.prune( |
| 536 | fn_inputs, fn_outputs, name, self.graph.structured_input_signature) |
| 537 | self._functions[name] = wrapped_function |
| 538 | return wrapped_function |
| 539 | |
| 540 | |
| 541 | @tf_export(v1=["wrap_function"]) |
no test coverage detected