a functional model is a model defined as a directed graph of layers three types of model exist subclassed model functional model and sequential a special case of functional a functional model can be instantiated by passing two arguments to init the first argument is the keras input objects that represent the inputs to the model the second argument specifies the output tensors that represent the outputs of this model both arguments can be a nested structure of tensors example inputs x1 keras inputshape10 name x1 x2 keras inputshape1 name x2 t keras layers dense1 activation relu inputs x1 outputs keras layers addt inputs x2 model keras modelinputs outputs a functional model constructed using the functional api can also include raw keras 3 ops example python inputs keras inputshape10 x keras layers dense1inputs outputs ops nn relux model keras modelinputs outputs a new functional model can also be using the intermediate tensors this enables you to quickly extract subcomponents of the model example python inputs keras inputshapenone none 3 processed keras layers randomcropwidth32 height32inputs conv keras layers conv2dfilters2 kernelsize3processed pooling keras layers globalaveragepooling2dconv feature keras layers dense10pooling fullmodel keras modelinputs feature backbone keras modelprocessed conv activations keras modelconv feature note that the backbone and activations models are not created with keras input objects but with the tensors that are originated from keras input objects under the hood the layers and weights will be shared across these models so that user can train the fullmodel and use backbone or activations to do feature extraction the inputs and outputs of the model can be nested structures of tensors as well and the created models are standard functional model that support all the existing api args inputs list of input tensors must be created via keras input or originated from keras input outputs list of output tensors name string optional name of the model trainable boolean optional if the model s variables should be trainable we will convert directly to the correct dtype per input unlike other layers we allow functional state to be mutable after build e g to attach a layer to a model that is not part of the functional dag add support for traning masking from function in the case that the graph is constructed with dict input tensors we will use the original dict key to map with the keys in the input data note that the model inputs is using tree flatten to process the input tensors which means the dict input tensors are ordered by their keys raise an warning if there are more input data comparing to input tensor flatten in the order inputs were passed during model construction otherwise both ref inputs and inputs will already be in same order add back metadata for backwards compatibility override input to retrieve the usedprovided constructor inputs symbolic only todo case where nestedinputs is a plain dict of inputs single input or listtuple of inputs the data may be passed as a dict keyed by input name subclassed networks are not serializable unless serialization is implemented by the of the subclassed network build a map from a layer unique name makenodekey to the index of the nodes that are saved in the config only nodes in networknodes are saved functional models start with a preexisting node linking their input to output i e we mark it to be saved serialize and save the layers in layerconfigs the node is relevant to the model add to filteredinboundnodes legacy format serialization used for h5 and savedmodel gather info about inputs and outputs instantiates a functional model from its config from getconfig args cls class of the model e g a custom subclass of model config output of getconfig for the original model instance customobjects optional dict of custom objects returns an instance of cls layer instances created during the graph reconstruction process dictionary mapping layer instances to node data that specifies a layer call it acts as a queue that maintains any unprocessed layer call until it becomes possible to process it i e until the input tensors to the call all exist add node to layer list arg layer layer object nodedata node data specifying layer call reconstruct node by linking to inbound layers args layer layer to process nodedata list of layer configs call layer on its inputs thus creating the node and building the layer if needed deserializes a layer then call it on appropriate inputs args layerdata layer config dict instantiate layer legacy format deserialization no module key used for h5 and savedmodel formats gather layer inputs we don t process nodes i e make layer calls on the fly because the inbound node may not yet exist in case of layer shared at different topological depths e g a model such as ababx first we create all layers and enqueue nodes to be processed then we process nodes in order of layer depth nodes that cannot yet be processed if the inbound node does not yet exist are reenqueued and the process is repeated until all nodes are processed process all nodes in layer if not yet processed process nodes in order if the node does not have all inbound layers available stop processing and continue later if not all nodes processed then store unprocessed nodes if all nodes processed remove the layer create lits of input and output tensors and return new class does not need to be serialized return args kwargs for calling the node layer if not nodedata return if isinstancenodedata list legacy case inputtensors for inputdata in nodedata inboundlayername inputdata0 inboundnodeindex inputdata1 inboundtensorindex inputdata2 if leninputdata 3 kwargs elif leninputdata 4 kwargs inputdata3 else raise valueerror cannot deserialize the model invalid config data inboundlayer createdlayersinboundlayername raise an error if the corresponding layer node has not yet been created if leninboundlayer inboundnodes inboundnodeindex raise indexerror layer node index out of bounds n finboundlayer inboundlayern inboundlayer inboundnodes finboundlayer inboundnodesn finboundnodeindex inboundnodeindex inboundnode inboundlayer inboundnodesinboundnodeindex inputtensors append inboundnode outputtensorsinboundtensorindex return unpacksingletoninputtensors kwargs args serializationlib deserializekerasobjectnodedataargs kwargs serializationlib deserializekerasobjectnodedatakwargs def convertrevivedtensorx if isinstancex backend kerastensor history x preserializationkerashistory if history is none return x layer createdlayers gethistory0 none if layer is none raise valueerrorfunknown layer history0 inboundnodeindex history1 inboundtensorindex history2 if lenlayer inboundnodes inboundnodeindex raise valueerror layer node index out of bounds n finboundlayer layern finboundlayer inboundnodes layer inboundnodesn finboundnodeindex inboundnodeindex inboundnode layer inboundnodesinboundnodeindex return inboundnode outputtensorsinboundtensorindex return x args tree mapstructureconvertrevivedtensor args kwargs tree mapstructureconvertrevivedtensor kwargs return args kwargs a functional model is a model defined as a directed graph of layers three types of model exist subclassed model functional model and sequential a special case of functional a functional model can be instantiated by passing two arguments to __init__ the first argument is the keras input objects that represent the inputs to the model the second argument specifies the output tensors that represent the outputs of this model both arguments can be a nested structure of tensors example inputs x1 keras input shape 10 name x1 x2 keras input shape 1 name x2 t keras layers dense 1 activation relu inputs x1 outputs keras layers add t inputs x2 model keras model inputs outputs a functional model constructed using the functional api can also include raw keras 3 ops example python inputs keras input shape 10 x keras layers dense 1 inputs outputs ops nn relu x model keras model inputs outputs a new functional model can also be using the intermediate tensors this enables you to quickly extract sub components of the model example python inputs keras input shape none none 3 processed keras layers randomcrop width 32 height 32 inputs conv keras layers conv2d filters 2 kernel_size 3 processed pooling keras layers globalaveragepooling2d conv feature keras layers dense 10 pooling full_model keras model inputs feature backbone keras model processed conv activations keras model conv feature note that the backbone and activations models are not created with keras input objects but with the tensors that are originated from keras input objects under the hood the layers and weights will be shared across these models so that user can train the full_model and use backbone or activations to do feature extraction the inputs and outputs of the model can be nested structures of tensors as well and the created models are standard functional model that support all the existing api args inputs list of input tensors must be created via keras input or originated from keras input outputs list of output tensors name string optional name of the model trainable boolean optional if the model s variables should be trainable we will convert directly to the correct dtype per input unlike other layers we allow functional state to be mutable after build e g to attach a layer to a model that is not part of the functional dag add support for traning masking from function in the case that the graph is constructed with dict input tensors we will use the original dict key to map with the keys in the input data note that the model inputs is using tree flatten to process the input tensors which means the dict input tensors are ordered by their keys raise an warning if there are more input data comparing to input tensor flatten in the order input s were passed during model construction otherwise both ref inputs and inputs will already be in same order add back metadata for backwards compatibility override input to retrieve the used provided constructor inputs symbolic only todo case where _nested_inputs is a plain dict of inputs single input or list tuple of inputs the data may be passed as a dict keyed by input name subclassed networks are not serializable unless serialization is implemented by the of the subclassed network build a map from a layer unique name make_node_key to the index of the nodes that are saved in the config only nodes in network_nodes are saved functional models start with a pre existing node linking their input to output i e we mark it to be saved serialize and save the layers in layer_configs from the earliest layers on the node is relevant to the model add to filtered_inbound_nodes legacy format serialization used for h5 and savedmodel gather info about inputs and outputs instantiates a functional model from its config from get_config args cls class of the model e g a custom subclass of model config output of get_config for the original model instance custom_objects optional dict of custom objects returns an instance of cls layer instances created during the graph reconstruction process dictionary mapping layer instances to node data that specifies a layer call it acts as a queue that maintains any unprocessed layer call until it becomes possible to process it i e until the input tensors to the call all exist add node to layer list arg layer layer object node_data node data specifying layer call reconstruct node by linking to inbound layers args layer layer to process node_data list of layer configs call layer on its inputs thus creating the node and building the layer if needed deserializes a layer then call it on appropriate inputs args layer_data layer config dict instantiate layer legacy format deserialization no module key used for h5 and savedmodel formats gather layer inputs we don t process nodes i e make layer calls on the fly because the inbound node may not yet exist in case of layer shared at different topological depths e g a model such as a b a b x first we create all layers and enqueue nodes to be processed then we process nodes in order of layer depth nodes that cannot yet be processed if the inbound node does not yet exist are re enqueued and the process is repeated until all nodes are processed process all nodes in layer if not yet processed process nodes in order if the node does not have all inbound layers available stop processing and continue later if not all nodes processed then store unprocessed nodes if all nodes processed remove the layer create lits of input and output tensors and return new class does not need to be serialized return args kwargs for calling the node layer legacy case raise an error if the corresponding layer node has not yet been created
import copy
import inspect
import warnings

import tree

from keras import backend
from keras import ops
from keras.backend.common import global_state
from keras.layers.input_spec import InputSpec
from keras.layers.layer import Layer
from keras.legacy.saving import saving_utils
from keras.legacy.saving import serialization as legacy_serialization
from keras.models.model import Model
from keras.ops.function import Function
from keras.ops.function import make_node_key
from keras.saving import serialization_lib
from keras.utils import tracking


class Functional(Function, Model):
    

    @tracking.no_automatic_dependency_tracking
    def __init__(self, inputs, outputs, name=None, **kwargs):
        if isinstance(inputs, dict):
            for k, v in inputs.items():
                if not isinstance(v, backend.KerasTensor):
                    raise ValueError(
                        "When providing `inputs` as a dict, all values in the "
                        f"dict must be KerasTensors. Received: inputs={inputs} "
                        f"including invalid value {v} of type {type(v)}"
                    )
                if k != v.name:
                    warnings.warn(
                        "When providing `inputs` as a dict, all keys in the "
                        "dict must match the names of the corresponding "
                        f"tensors. Received key '{k}' mapping to value {v} "
                        f"which has name '{v.name}'. Change the tensor name to "
                        f"'{k}' (via `Input(..., name='{k}')`)"
                    )
        elif isinstance(inputs, (list, tuple)):
            for x in inputs:
                if not isinstance(x, backend.KerasTensor):
                    raise ValueError(
                        "When providing `inputs` as a list/tuple, all values "
                        f"in the list/tuple must be KerasTensors. Received: "
                        f"inputs={inputs} including invalid value {x} of type "
                        f"{type(x)}"
                    )
        elif not isinstance(inputs, backend.KerasTensor):
            raise ValueError(
                f"Unrecognized type for `inputs`: {inputs} "
                f"(of type {type(inputs)})"
            )
        if isinstance(outputs, dict):
            for k, v in outputs.items():
                if not isinstance(v, backend.KerasTensor):
                    raise ValueError(
                        "When providing `outputs` as a dict, all values in the "
                        f"dict must be KerasTensors. Received: "
                        f"outputs={outputs} including invalid value {v} of "
                        f"type {type(v)}"
                    )
        elif isinstance(outputs, (list, tuple)):
            for x in outputs:
                if not isinstance(x, backend.KerasTensor):
                    raise ValueError(
                        "When providing `outputs` as a list/tuple, all values "
                        f"in the list/tuple must be KerasTensors. Received: "
                        f"outputs={outputs} including invalid value {x} of "
                        f"type {type(x)}"
                    )
        elif not isinstance(outputs, backend.KerasTensor):
            raise ValueError(
                f"Unrecognized type for `outputs`: {outputs} "
                f"(of type {type(outputs)})"
            )

        trainable = kwargs.pop("trainable", None)

        Function.__init__(self, inputs, outputs, name=name, **kwargs)

        if trainable is not None:
            self.trainable = trainable

        self._layers = self.layers
        self.build(None)
        
        self._convert_input_args = False
        self._allow_non_tensor_positional_args = True
        output_layers = [x._keras_history[0] for x in self.outputs]
        self.output_names = [x.name for x in output_layers]

    def _lock_state(self):
        
        
        
        pass

    @property
    def layers(self):
        layers = []
        for operation in self._operations:
            if isinstance(operation, Layer):
                layers.append(operation)
        return layers

    def call(self, inputs, training=None, mask=None):
        
        inputs = self._standardize_inputs(inputs)
        if mask is None:
            masks = [None] * len(inputs)
        else:
            masks = self._flatten_to_reference_inputs(mask)
            for x, mask in zip(inputs, masks):
                if mask is not None:
                    x._keras_mask = mask
        outputs = self._run_through_graph(
            inputs, operation_fn=lambda op: operation_fn(op, training=training)
        )
        return unpack_singleton(outputs)

    def compute_output_spec(self, inputs, training=None, mask=None):
        
        return super().compute_output_spec(inputs)

    def build(self, input_shape):
        self.built = True

    @property
    def input_shape(self):
        input_shapes = tree.map_structure(lambda x: x.shape, self.inputs)
        if isinstance(input_shapes, list) and len(input_shapes) == 1:
            return input_shapes[0]
        return input_shapes

    @property
    def output_shape(self):
        output_shapes = tree.map_structure(lambda x: x.shape, self.outputs)
        if isinstance(output_shapes, list) and len(output_shapes) == 1:
            return output_shapes[0]
        return output_shapes

    def _assert_input_compatibility(self, *args):
        return super(Model, self)._assert_input_compatibility(*args)

    def _flatten_to_reference_inputs(self, inputs, allow_extra_keys=True):
        if isinstance(inputs, dict):
            ref_inputs = self._inputs_struct
            if not tree.is_nested(ref_inputs):
                ref_inputs = [self._inputs_struct]
            if isinstance(ref_inputs, dict):
                
                
                
                
                
                ref_input_names = sorted(ref_inputs.keys())
            else:
                ref_input_names = [
                    inp._keras_history.operation.name for inp in ref_inputs
                ]
            
            
            if not allow_extra_keys and len(inputs) > len(ref_input_names):
                warnings.warn(
                    "Input dict contained keys {} which did not match any "
                    "model input. They will be ignored by the model.".format(
                        [n for n in inputs.keys() if n not in ref_input_names]
                    ),
                    stacklevel=2,
                )
            
            
            return [inputs[n] for n in ref_input_names]
        
        return tree.flatten(inputs)

    def _convert_inputs_to_tensors(self, flat_inputs):
        converted = []
        for x, input in zip(flat_inputs, self._inputs):
            converted.append(
                ops.convert_to_tensor(x, dtype=input.dtype, sparse=input.sparse)
            )
        return converted

    def _adjust_input_rank(self, flat_inputs):
        flat_ref_shapes = [x.shape for x in self._inputs]
        adjusted = []
        for x, ref_shape in zip(flat_inputs, flat_ref_shapes):
            x_rank = len(x.shape)
            ref_rank = len(ref_shape)
            if x_rank == ref_rank:
                adjusted.append(x)
                continue
            if x_rank == ref_rank + 1:
                if x.shape[-1] == 1:
                    adjusted.append(ops.squeeze(x, axis=-1))
                    continue
            if x_rank == ref_rank - 1:
                if ref_shape[-1] == 1:
                    adjusted.append(ops.expand_dims(x, axis=-1))
                    continue
            raise ValueError(
                f"Invalid input shape for input {x}. Expected shape "
                f"{ref_shape}, but input has incompatible shape {x.shape}"
            )
        
        for i in range(len(flat_inputs)):
            if hasattr(flat_inputs[i], "_keras_history"):
                adjusted[i]._keras_history = flat_inputs[i]._keras_history
            if hasattr(flat_inputs[i], "_keras_mask"):
                adjusted[i]._keras_mask = flat_inputs[i]._keras_mask
        return adjusted

    def _standardize_inputs(self, inputs):
        flat_inputs = self._flatten_to_reference_inputs(inputs)
        flat_inputs = self._convert_inputs_to_tensors(flat_inputs)
        return self._adjust_input_rank(flat_inputs)

    @property
    def input(self):
        
        
        
        return self._inputs_struct

    @property
    def output(self):
        return self._outputs_struct

    def add_loss(self, loss):
        
        raise NotImplementedError

    @property
    def input_spec(self):
        if hasattr(self, "_manual_input_spec"):
            return self._manual_input_spec

        def shape_with_no_batch_size(x):
            x = list(x)
            if x:
                x[0] = None
            return tuple(x)

        if isinstance(self._inputs_struct, dict):
            
            names = sorted(self._inputs_struct.keys())
            return [
                InputSpec(
                    shape=shape_with_no_batch_size(
                        self._inputs_struct[name].shape
                    ),
                    allow_last_axis_squeeze=True,
                    name=name,
                )
                for name in names
            ]
        else:
            
            
            return [
                InputSpec(
                    shape=shape_with_no_batch_size(x.shape),
                    allow_last_axis_squeeze=True,
                    name=x._keras_history[0].name,
                )
                for x in self._inputs
            ]

    @input_spec.setter
    def input_spec(self, value):
        self._manual_input_spec = value

    def get_config(self):
        if not functional_like_constructor(self.__class__):
            
            
            
            return Model.get_config(self)

        config = {
            "name": self.name,
            "trainable": self.trainable,
        }
        
        
        
        node_reindexing_map = {}
        for operation in self.operations:
            if issubclass(operation.__class__, Functional):
                
                
                kept_nodes = 1
            else:
                kept_nodes = 0
            for original_node_index, node in enumerate(
                operation._inbound_nodes
            ):
                node_key = make_node_key(operation, original_node_index)
                if node_key in self._nodes:
                    
                    node_reindexing_map[node_key] = kept_nodes
                    kept_nodes += 1

        
        layer_configs = []
        for operation in self.operations:  
            filtered_inbound_nodes = []
            for original_node_index, node in enumerate(
                operation._inbound_nodes
            ):
                node_key = make_node_key(operation, original_node_index)
                if node_key in self._nodes:
                    
                    
                    node_data = serialize_node(node, node_reindexing_map)
                    if node_data is not None:
                        filtered_inbound_nodes.append(node_data)

            serialize_obj_fn = serialization_lib.serialize_keras_object
            if global_state.get_global_attribute("use_legacy_config", False):
                
                serialize_obj_fn = legacy_serialization.serialize_keras_object
            layer_config = serialize_obj_fn(operation)
            layer_config["name"] = operation.name
            layer_config["inbound_nodes"] = filtered_inbound_nodes
            layer_configs.append(layer_config)
        config["layers"] = layer_configs

        
        def get_tensor_config(tensor):
            operation = tensor._keras_history[0]
            node_index = tensor._keras_history[1]
            tensor_index = tensor._keras_history[2]
            node_key = make_node_key(operation, node_index)
            assert node_key in self._nodes
            new_node_index = node_reindexing_map[node_key]
            return [operation.name, new_node_index, tensor_index]

        def map_tensors(tensors):
            if isinstance(tensors, dict):
                return {k: get_tensor_config(v) for k, v in tensors.items()}
            if isinstance(tensors, (list, tuple)):
                return [get_tensor_config(v) for v in tensors]
            else:
                return [get_tensor_config(tensors)]

        config["input_layers"] = map_tensors(self._inputs_struct)
        config["output_layers"] = map_tensors(self._outputs_struct)
        return copy.deepcopy(config)


def functional_from_config(cls, config, custom_objects=None):
    
    
    
    created_layers = {}

    
    
    
    
    
    unprocessed_nodes = {}

    def add_unprocessed_node(layer, node_data):
        
        if layer not in unprocessed_nodes:
            unprocessed_nodes[layer] = [node_data]
        else:
            unprocessed_nodes[layer].append(node_data)

    def process_node(layer, node_data):
        
        args, kwargs = deserialize_node(node_data, created_layers)
        
        
        layer(*args, **kwargs)

    def process_layer(layer_data):
        
        layer_name = layer_data["name"]

        
        if "module" not in layer_data:
            
            
            layer = saving_utils.model_from_config(
                layer_data, custom_objects=custom_objects
            )
        else:
            layer = serialization_lib.deserialize_keras_object(
                layer_data, custom_objects=custom_objects
            )
        created_layers[layer_name] = layer

        
        inbound_nodes_data = layer_data["inbound_nodes"]
        for node_data in inbound_nodes_data:
            
            
            
            
            add_unprocessed_node(layer, node_data)

    
    for layer_data in config["layers"]:
        process_layer(layer_data)

    
    
    
    
    while unprocessed_nodes:
        for layer_data in config["layers"]:
            layer = created_layers[layer_data["name"]]

            
            if layer in unprocessed_nodes:
                node_data_list = unprocessed_nodes[layer]

                
                node_index = 0
                while node_index < len(node_data_list):
                    node_data = node_data_list[node_index]
                    try:
                        process_node(layer, node_data)

                    
                    
                    except IndexError:
                        break

                    node_index += 1

                
                if node_index < len(node_data_list):
                    unprocessed_nodes[layer] = node_data_list[node_index:]
                
                else:
                    del unprocessed_nodes[layer]

    
    name = config.get("name")
    trainable = config.get("trainable")

    def get_tensor(layer_name, node_index, tensor_index):
        assert layer_name in created_layers
        layer = created_layers[layer_name]
        layer_output_tensors = layer._inbound_nodes[node_index].output_tensors
        return layer_output_tensors[tensor_index]

    def map_tensors(tensors):
        if isinstance(tensors, dict):
            return {k: get_tensor(*v) for k, v in tensors.items()}
        else:
            return [get_tensor(*v) for v in tensors]

    input_tensors = map_tensors(config["input_layers"])
    output_tensors = map_tensors(config["output_layers"])
    return cls(
        inputs=input_tensors,
        outputs=output_tensors,
        name=name,
        trainable=trainable,
    )


def operation_fn(operation, training):
    def call(*args, **kwargs):
        if (
            hasattr(operation, "_call_has_training_arg")
            and operation._call_has_training_arg
            and training is not None
        ):
            kwargs["training"] = training
        return operation(*args, **kwargs)

    return call


def functional_like_constructor(cls):
    init_args = inspect.getfullargspec(cls.__init__).args[1:]
    functional_init_args = inspect.getfullargspec(Functional.__init__).args[1:]
    if init_args == functional_init_args:
        return True
    return False


def unpack_singleton(x):
    if isinstance(x, (list, tuple)) and len(x) == 1:
        return x[0]
    return x


def serialize_node(node, node_reindexing_map):
    if not node.input_tensors:
        
        return

    args = node.arguments.args
    kwargs = node.arguments.kwargs
    return {
        "args": serialization_lib.serialize_keras_object(args),
        "kwargs": serialization_lib.serialize_keras_object(kwargs),
    }


def deserialize_node(node_data, created_layers):
    
    if not node_data:
        return [], {}

    if isinstance(node_data, list):
        
        input_tensors = []
        for input_data in node_data:
            inbound_layer_name = input_data[0]
            inbound_node_index = input_data[1]
            inbound_tensor_index = input_data[2]
            if len(input_data) == 3:
                kwargs = {}
            elif len(input_data) == 4:
                kwargs = input_data[3]
            else:
                raise ValueError(
                    "Cannot deserialize the model (invalid config data?)"
                )
            inbound_layer = created_layers[inbound_layer_name]

            
            
            if len(inbound_layer._inbound_nodes) <= inbound_node_index:
                raise IndexError(
                    "Layer node index out of bounds.\n"
                    f"inbound_layer = {inbound_layer}\n"
                    "inbound_layer._inbound_nodes = "
                    f"{inbound_layer._inbound_nodes}\n"
                    f"inbound_node_index = {inbound_node_index}"
                )
            inbound_node = inbound_layer._inbound_nodes[inbound_node_index]
            input_tensors.append(
                inbound_node.output_tensors[inbound_tensor_index]
            )
        return [unpack_singleton(input_tensors)], kwargs

    args = serialization_lib.deserialize_keras_object(node_data["args"])
    kwargs = serialization_lib.deserialize_keras_object(node_data["kwargs"])

    def convert_revived_tensor(x):
        if isinstance(x, backend.KerasTensor):
            history = x._pre_serialization_keras_history
            if history is None:
                return x
            layer = created_layers.get(history[0], None)
            if layer is None:
                raise ValueError(f"Unknown layer: {history[0]}")
            inbound_node_index = history[1]
            inbound_tensor_index = history[2]
            if len(layer._inbound_nodes) <= inbound_node_index:
                raise ValueError(
                    "Layer node index out of bounds.\n"
                    f"inbound_layer = {layer}\n"
                    f"inbound_layer._inbound_nodes = {layer._inbound_nodes}\n"
                    f"inbound_node_index = {inbound_node_index}"
                )
            inbound_node = layer._inbound_nodes[inbound_node_index]
            return inbound_node.output_tensors[inbound_tensor_index]
        return x

    args = tree.map_structure(convert_revived_tensor, args)
    kwargs = tree.map_structure(convert_revived_tensor, kwargs)
    return args, kwargs