Class WrappedFunction

tensorflow/python/eager/wrap_function.py:218–361

Wraps a tf V1 piece of code in a function.

# TODO(allenl): make this trackable
class WrappedFunction(function.ConcreteFunction):
  """Wraps a tf V1 piece of code in a function."""

  def __init__(self, fn_graph, variable_holder, attrs=None, signature=None):
    self._variable_holder = variable_holder
    _lift_unlifted_variables(fn_graph, variable_holder)
    # We call __init__ after lifting variables so that the function's signature
    # properly reflects the new captured inputs.
    for f in fn_graph.as_graph_def().library.function:
      context.context().add_function_def(f)
    super(WrappedFunction, self).__init__(
        fn_graph, attrs=attrs, signature=signature)

  def prune(self, feeds, fetches, name=None, input_signature=None):
    """Extract a subgraph of this function's underlying graph.

    Wraps the subgraph in a new `WrappedFunction` object.

    Args:
      feeds: Input tensors to the subgraph to extract, as `Tensor` objects.
      fetches: Possibly-nested Python data structure containing information
        about outputs of the target subgraph. Each entry can either be a
        `Tensor` object (for data outputs), an `Operation` object (for control
        outputs), or a `TensorInfo` proto. Any additional shape/dtype
        information provided in a `TensorInfo` and not present in the original
        graph will be added to the returned subgraph.
      name: (optional) Name to give to the underlying `FuncGraph` of the
        returned object. If no name is provided, the graph's name will be
        `"pruned"`.
      input_signature: (optional) possibly-nested Python data structure
        containing `TensorSpec` objects, with which to populate the returned
        functions's `FuncGraph`'s `structured_input_signature` field.

    Returns:
      A new `WrappedFunction` object containing a copy of the portion of this
        object's graph that goes from `feeds` to `fetches`.
    """
    # TODO(b/129646028): Add support for CompositeTensors.
    name = name or "pruned"
    flat_feeds = nest.flatten(feeds, expand_composites=True)
    flat_feeds = [self.graph.as_graph_element(t) for t in flat_feeds]
    for f in flat_feeds:
      if not isinstance(f, ops.Tensor):
        raise ValueError("Feeds must be tensors.")

    # Ignoring all feeds that are captures allows prune to be called
    # using wrapped_func.inputs even when it uses variables
    internal_captures = object_identity.ObjectIdentitySet(
        self.graph.internal_captures)
    flat_feeds = [f for f in flat_feeds if f not in internal_captures]

    operation_fetches = []
    tensor_fetches = []
    tensor_infos = []

    def _fetch_preprocesing_callback(fetch):
      """Extract out lists of ops, tensors, and tensor type info.