class that encapsulates a computation graph of keras operations you can use a function to capture the computation graph linking some input tensors to some output tensors and reapply the same computation on new inputs a function is similar to a functional model with the difference that it is stateless it does not track state variables and does not implement the layer api example python input1 keras kerastensorshapenone 2 3 input2 keras kerastensorshapenone 2 3 x input1 input2 output keras ops sigmoidx fn keras functioninputsinput1 input2 outputsoutput input1val np random random4 2 3 input2val np random random4 2 3 outputval fninput1val input2val args inputs kerastensor instance or nested structured of kerastensor instances outputs kerastensor instance or nested structured of kerastensor instances they should be computable given only the values of inputs name string the name of the function temporary work around for https github comkerasteamkerasissues931 this stop tensorflow from wrapping tf function output in a dictwrapper object check if input shapes are identical to ref input shapes if so take a shortcut no luck take the long road through the graph original keras used a cache to avoid recomputing all this when known input shapes where seen again perhaps a good idea to bring that back computes output tensors for new inputs self assertinputcompatibilityinputs return self runthroughgraphinputs operationfnlambda op op def runthroughgraphself inputs operationfn inputs tree flatteninputs dictionary mapping reference tensors to computed tensors tensordict for x y in zipself inputs inputs tensordictidx y nodesbydepth self nodesbydepth depthkeys listnodesbydepth keys depthkeys sortreversetrue for depth in depthkeys nodes nodesbydepthdepth for node in nodes if not node operation or node isinput continue input tensors already exist if anyidx not in tensordict for x in node inputtensors continue node is not computable try skipping args kwargs node arguments fillintensordict outputs operationfnnode operationargs kwargs update tensordict for x y in zipnode outputs tree flattenoutputs tensordictidx y outputtensors for x in self outputs outputtensors appendtensordictidx return packsequenceasself outputsstruct outputtensors def assertinputcompatibilityself inputs try tree assertsamestructure inputs self inputsstruct checktypesfalse except valueerror raise valueerror function was called with an invalid input structure fexpected input structure self inputsstructn freceived input structure inputs for x xref in ziptree flatteninputs self inputs if lenx shape lenxref shape raise valueerror fself class name was passed fincompatible inputs for input xref name fexpected shape xref shape but received finstead a tensor with shape x shape for dim refdim in zipx shape xref shape if refdim is not none and dim is not none if dim refdim raise valueerror fself class name was passed fincompatible inputs for input xref name fexpected shape xref shape but received finstead a tensor with shape x shape def makenodekeyop nodeindex return stridop ib strnodeindex def mapgraphinputs outputs depth is number of operations between output node and the node nodes are ordered from inputs outputs nodesindecreasingdepth operationindices buildmapoutputs networknodes makenodekeynode operation node operation inboundnodes indexnode for node in nodesindecreasingdepth nodesdepths dict node depth value operationsdepths dict operation depth value for node in reversednodesindecreasingdepth if the depth is not set the node has no outbound nodes depth 0 depth nodesdepths setdefaultnode 0 update the depth of the corresponding operation previousdepth operationsdepths getnode operation 0 if we ve seen this operation before at a higher depth we should use that depth instead of the node depth this is necessary for shared operations that have inputs at different depth levels in the graph depth maxdepth previousdepth operationsdepthsnode operation depth nodesdepthsnode depth update the depth of inbound nodes the depth of a node is the max of the depths of all nodes it is connected to 1 for nodedep in node parentnodes previousdepth nodesdepths getnodedep 0 nodesdepthsnodedep maxdepth 1 previousdepth handle inputs that are not connected to outputs we do not error out here because the inputs may be used to compute losses and metrics for inputt in inputs inputoperation inputt kerashistory0 if inputoperation and inputoperation not in operationsdepths operationsdepthsinputoperation 0 operationindicesinputoperation 1 nodesdepthsinputoperation inboundnodes0 0 networknodes addmakenodekeyinputoperation 0 build a dict depth list of nodes with this depth nodesbydepth collections defaultdictlist for node depth in nodesdepths items nodesbydepthdepth appendnode build a dict depth list of operations with this depth operationsbydepth collections defaultdictlist for operation depth in operationsdepths items operationsbydepthdepth appendoperation get sorted list of operation depths depthkeys listoperationsbydepth keys depthkeys sortreversetrue set self operations ordered by depth operations for depth in depthkeys operationsfordepth operationsbydepthdepth network operations needs to have a deterministic order here we order them by traversal order operationsfordepth sortkeylambda x operationindicesx operations extendoperationsfordepth get sorted list of node depths depthkeys listnodesbydepth keys depthkeys sortreversetrue check that all tensors required are computable computabletensors all tensors in the graph that can be computed from the inputs provided computabletensors set for x in inputs computabletensors addx operationswithcompleteinput to provide a better error msg for depth in depthkeys for node in nodesbydepthdepth for x in tree flattennode inputtensors if x not in computabletensors operation node operation raise valueerror graph disconnected cannot find parent for ftensor x at operation operation the following previous operations were accessed fwithout issue operationswithcompleteinput operationswithcompleteinput appendoperation name for x in tree flattennode outputs computabletensors addx ensure name unicity which will be crucial for serialization since serialized nodes refer to operations by their name allnames operation name for operation in operations for name in allnames if allnames countname 1 raise valueerror f the name name is used allnames countname times in the model all operation names should be unique return networknodes nodesbydepth operations operationsbydepth def buildmapoutputs finishednodes set nodesinprogress set nodesindecreasingdepth nodes from inputs outputs operationindices operation in traversal order for output in tree flattenoutputs buildmaphelper output finishednodes nodesinprogress nodesindecreasingdepth operationindices return nodesindecreasingdepth operationindices def buildmaphelper tensor finishednodes nodesinprogress nodesindecreasingdepth operationindices don t repeat work for shared subgraphs prevent cycles store the traversal order for operation sorting propagate to all previous tensors connected to this node class that encapsulates a computation graph of keras operations you can use a function to capture the computation graph linking some input tensors to some output tensors and reapply the same computation on new inputs a function is similar to a functional model with the difference that it is stateless it does not track state variables and does not implement the layer api example python input_1 keras kerastensor shape none 2 3 input_2 keras kerastensor shape none 2 3 x input_1 input_2 output keras ops sigmoid x fn keras function inputs input_1 input_2 outputs output input_1_val np random random 4 2 3 input_2_val np random random 4 2 3 output_val fn input_1_val input_2_val args inputs kerastensor instance or nested structured of kerastensor instances outputs kerastensor instance or nested structured of kerastensor instances they should be computable given only the values of inputs name string the name of the function temporary work around for https github com keras team keras issues 931 this stop tensorflow from wrapping tf function output in a _dictwrapper object check if input shapes are identical to ref input shapes if so take a shortcut no luck take the long road through the graph original keras used a cache to avoid recomputing all this when known input shapes where seen again perhaps a good idea to bring that back computes output tensors for new inputs execute the graph at each node we compute outputs via operation_fn node operation args kwargs dictionary mapping reference tensors to computed tensors input tensors already exist node is not computable try skipping update tensor_dict validates a graph s topology and gather its operations and nodes args inputs list of input tensors outputs list of outputs tensors returns a tuple nodes nodes_by_depth operations operations_by_depth network_nodes dict mapping unique node keys to the node instances nodes_by_depth dict mapping ints depth to lists of node instances operations list of operation instances operations_by_depth dict mapping ints depth to lists of operation instances depth is number of operations between output node and the node nodes are ordered from inputs outputs dict node depth value dict operation depth value if the depth is not set the node has no outbound nodes depth 0 update the depth of the corresponding operation if we ve seen this operation before at a higher depth we should use that depth instead of the node depth this is necessary for shared operations that have inputs at different depth levels in the graph update the depth of inbound nodes the depth of a node is the max of the depths of all nodes it is connected to 1 handle inputs that are not connected to outputs we do not error out here because the inputs may be used to compute losses and metrics build a dict depth list of nodes with this depth build a dict depth list of operations with this depth get sorted list of operation depths set self operations ordered by depth network operations needs to have a deterministic order here we order them by traversal order get sorted list of node depths check that all tensors required are computable computable_tensors all tensors in the graph that can be computed from the inputs provided to provide a better error msg ensure name unicity which will be crucial for serialization since serialized nodes refer to operations by their name topologically sort nodes in order from inputs to outputs it uses a depth first search to topologically sort nodes that appear in the _keras_history connectivity metadata of outputs args outputs the output tensors whose _keras_history metadata should be walked this may be an arbitrary nested structure returns a tuple like ordered_nodes operation_to_first_traversal_index ordered_nodes list of nodes appearing in the keras history topologically sorted from original inputs to the outputs if outputs have different sets of ancestors the inputs to one output may appear after a different output operation_to_first_traversal_index a dict mapping operation to the traversal index in the dfs where it is seen note if a operation is shared by several nodes the dict will onlystore the index corresponding to the first time the operation seen nodes from inputs outputs operation in traversal order recursive helper for _build_map don t repeat work for shared subgraphs prevent cycles store the traversal order for operation sorting propagate to all previous tensors connected to this node
import collections

import tree

from keras.api_export import keras_export
from keras.backend import KerasTensor
from keras.backend.config import backend
from keras.ops.operation import Operation
from keras.utils.nest import pack_sequence_as


@keras_export("keras.Function")
class Function(Operation):
    

    def __init__(self, inputs, outputs, name=None):
        super().__init__(name=name)

        if backend() == "tensorflow":
            
            
            
            
            _self_setattr_tracking = getattr(
                self, "_self_setattr_tracking", True
            )
            self._self_setattr_tracking = False
        self._inputs_struct = tree.map_structure(lambda x: x, inputs)
        self._outputs_struct = tree.map_structure(lambda x: x, outputs)
        self._inputs = tree.flatten(inputs)
        self._outputs = tree.flatten(outputs)
        if not self._inputs:
            raise ValueError(
                "`inputs` argument cannot be empty. Received:\n"
                f"inputs={inputs}\n"
                f"outputs={outputs}"
            )
        if not self._outputs:
            raise ValueError(
                "`outputs` argument cannot be empty. Received:\n"
                f"inputs={inputs}\n"
                f"outputs={outputs}"
            )

        if backend() == "tensorflow":
            self._self_setattr_tracking = _self_setattr_tracking

        (nodes, nodes_by_depth, operations, operations_by_depth) = map_graph(
            self._inputs, self._outputs
        )
        self._nodes = nodes
        self._nodes_by_depth = nodes_by_depth
        self._operations = operations
        self._operations_by_depth = operations_by_depth

    @property
    def operations(self):
        return self._operations[:]

    @property
    def inputs(self):
        return self._inputs

    @property
    def outputs(self):
        return self._outputs

    def compute_output_spec(self, inputs):
        self._assert_input_compatibility(inputs)
        
        
        shortcut = True
        for x, x_ref in zip(tree.flatten(inputs), self._inputs):
            if x.shape != x_ref.shape:
                shortcut = False
                break
        if shortcut:
            return tree.map_structure(
                lambda x: KerasTensor(shape=x.shape, dtype=x.dtype),
                self._outputs_struct,
            )
        
        
        
        
        return self._run_through_graph(
            inputs, operation_fn=lambda op: op.compute_output_spec
        )

    def call(self, inputs):
        
        self._assert_input_compatibility(inputs)
        return self._run_through_graph(inputs, operation_fn=lambda op: op)

    def _run_through_graph(self, inputs, operation_fn):
        
        inputs = tree.flatten(inputs)

        
        tensor_dict = {}
        for x, y in zip(self.inputs, inputs):
            tensor_dict[id(x)] = y

        nodes_by_depth = self._nodes_by_depth
        depth_keys = list(nodes_by_depth.keys())
        depth_keys.sort(reverse=True)

        for depth in depth_keys:
            nodes = nodes_by_depth[depth]
            for node in nodes:
                if not node.operation or node.is_input:
                    continue  

                if any(id(x) not in tensor_dict for x in node.input_tensors):
                    continue  

                args, kwargs = node.arguments.fill_in(tensor_dict)
                outputs = operation_fn(node.operation)(*args, **kwargs)

                
                for x, y in zip(node.outputs, tree.flatten(outputs)):
                    tensor_dict[id(x)] = y

        output_tensors = []
        for x in self.outputs:
            output_tensors.append(tensor_dict[id(x)])

        return pack_sequence_as(self._outputs_struct, output_tensors)

    def _assert_input_compatibility(self, inputs):
        try:
            tree.assert_same_structure(
                inputs, self._inputs_struct, check_types=False
            )
        except ValueError:
            raise ValueError(
                "Function was called with an invalid input structure. "
                f"Expected input structure: {self._inputs_struct}\n"
                f"Received input structure: {inputs}"
            )
        for x, x_ref in zip(tree.flatten(inputs), self._inputs):
            if len(x.shape) != len(x_ref.shape):
                raise ValueError(
                    f"{self.__class__.__name__} was passed "
                    f"incompatible inputs. For input '{x_ref.name}', "
                    f"expected shape {x_ref.shape}, but received "
                    f"instead a tensor with shape {x.shape}."
                )
            for dim, ref_dim in zip(x.shape, x_ref.shape):
                if ref_dim is not None and dim is not None:
                    if dim != ref_dim:
                        raise ValueError(
                            f"{self.__class__.__name__} was passed "
                            f"incompatible inputs. For input '{x_ref.name}', "
                            f"expected shape {x_ref.shape}, but received "
                            f"instead a tensor with shape {x.shape}."
                        )


def make_node_key(op, node_index):
    return str(id(op)) + "_ib-" + str(node_index)


def map_graph(inputs, outputs):
    
    
    
    nodes_in_decreasing_depth, operation_indices = _build_map(outputs)
    network_nodes = {
        make_node_key(node.operation, node.operation._inbound_nodes.index(node))
        for node in nodes_in_decreasing_depth
    }

    nodes_depths = {}  
    operations_depths = {}  

    for node in reversed(nodes_in_decreasing_depth):
        
        depth = nodes_depths.setdefault(node, 0)

        
        previous_depth = operations_depths.get(node.operation, 0)
        
        
        
        
        depth = max(depth, previous_depth)
        operations_depths[node.operation] = depth
        nodes_depths[node] = depth

        
        
        
        for node_dep in node.parent_nodes:
            previous_depth = nodes_depths.get(node_dep, 0)
            nodes_depths[node_dep] = max(depth + 1, previous_depth)

    
    
    
    for input_t in inputs:
        input_operation = input_t._keras_history[0]
        if input_operation and input_operation not in operations_depths:
            operations_depths[input_operation] = 0
            operation_indices[input_operation] = -1
            nodes_depths[input_operation._inbound_nodes[0]] = 0
            network_nodes.add(make_node_key(input_operation, 0))

    
    nodes_by_depth = collections.defaultdict(list)
    for node, depth in nodes_depths.items():
        nodes_by_depth[depth].append(node)

    
    operations_by_depth = collections.defaultdict(list)
    for operation, depth in operations_depths.items():
        operations_by_depth[depth].append(operation)

    
    depth_keys = list(operations_by_depth.keys())
    depth_keys.sort(reverse=True)

    
    operations = []
    for depth in depth_keys:
        operations_for_depth = operations_by_depth[depth]
        
        
        operations_for_depth.sort(key=lambda x: operation_indices[x])
        operations.extend(operations_for_depth)

    
    depth_keys = list(nodes_by_depth.keys())
    depth_keys.sort(reverse=True)

    
    
    
    computable_tensors = set()
    for x in inputs:
        computable_tensors.add(x)

    operations_with_complete_input = []  
    for depth in depth_keys:
        for node in nodes_by_depth[depth]:
            for x in tree.flatten(node.input_tensors):
                if x not in computable_tensors:
                    operation = node.operation
                    raise ValueError(
                        "Graph disconnected: cannot find parent for "
                        f"tensor {x} at operation '{operation}'. "
                        "The following previous operations were accessed "
                        f"without issue: {operations_with_complete_input}"
                    )
                operations_with_complete_input.append(operation.name)

            for x in tree.flatten(node.outputs):
                computable_tensors.add(x)

    
    
    all_names = [operation.name for operation in operations]
    for name in all_names:
        if all_names.count(name) != 1:
            raise ValueError(
                f'The name "{name}" is used {all_names.count(name)} '
                "times in the model. All operation names should be unique."
            )
    return network_nodes, nodes_by_depth, operations, operations_by_depth


def _build_map(outputs):
    
    finished_nodes = set()
    nodes_in_progress = set()
    nodes_in_decreasing_depth = []  
    operation_indices = {}  
    for output in tree.flatten(outputs):
        _build_map_helper(
            output,
            finished_nodes,
            nodes_in_progress,
            nodes_in_decreasing_depth,
            operation_indices,
        )
    return nodes_in_decreasing_depth, operation_indices


def _build_map_helper(
    tensor,
    finished_nodes,
    nodes_in_progress,
    nodes_in_decreasing_depth,
    operation_indices,
):
    
    (
        operation,
        node_index,
        _,
    ) = tensor._keras_history
    if not operation:
        return

    node = operation._inbound_nodes[node_index]

    
    if node in finished_nodes:
        return

    
    if node in nodes_in_progress:
        raise ValueError(
            f"Tensor {tensor} from operation '{operation.name}' is part of a "
            "cycle."
        )

    
    if operation not in operation_indices:
        operation_indices[operation] = len(operation_indices)

    
    nodes_in_progress.add(node)
    if not node.is_input:
        for tensor in node.input_tensors:
            _build_map_helper(
                tensor,
                finished_nodes,
                nodes_in_progress,
                nodes_in_decreasing_depth,
                operation_indices,
            )

    finished_nodes.add(node)
    nodes_in_progress.remove(node)
    nodes_in_decreasing_depth.append(node)