library for exporting inferenceonly keras modelslayers from keras import backend from keras apiexport import kerasexport from keras layers import layer from keras models import functional from keras models import sequential from keras utils import ioutils from keras utils moduleutils import tensorflow as tf kerasexportkeras export exportarchive class exportarchive elsewhere we can reload the artifact and serve it the endpoint we added is available as a method def initself self endpointnames self endpointsignatures self tensorflowversion tf version self tftrackable tf internal tracking autotrackable self tftrackable variables self tftrackable trainablevariables self tftrackable nontrainablevariables if backend backend not in tensorflow jax raise notimplementederror the export api is only compatible with jax and tf backends todonkovela make jax version checking programatic if backend backend jax from jax import version as jaxv if jaxv 0 4 15 raise valueerror the export api is only compatible with jax version 0 4 15 fand prior your jax version jaxv property def variablesself return self tftrackable variables property def trainablevariablesself return self tftrackable trainablevariables property def nontrainablevariablesself return self tftrackable nontrainablevariables def trackself resource layers in tracked are not part of the trackables that get saved because we re creating the attribute in a noautomaticdependencytracking scope variables in the lists below are actually part of the trackables that get saved because the lists are created in init register a new serving endpoint arguments name str name of the endpoint fn a function it should only leverage resources e g tf variable objects or tf lookup statichashtable objects that are available on the modelslayers tracked by the exportarchive you can call trackmodel to track a new model the shape and dtype of the inputs to the function must be known for that purpose you can either 1 make sure that fn is a tf function that has been called at least once or 2 provide an inputsignature argument that specifies the shape and dtype of the inputs see below inputsignature used to specify the shape and dtype of the inputs to fn list of tf tensorspec objects one per positional input argument of fn nested arguments are allowed see below for an example showing a functional model with 2 input arguments example adding an endpoint using the inputsignature argument when the model has a single input argument python exportarchive exportarchive exportarchive trackmodel exportarchive addendpoint nameserve fnmodel call inputsignaturetf tensorspecshapenone 3 dtypetf float32 adding an endpoint using the inputsignature argument when the model has two positional input arguments python exportarchive exportarchive exportarchive trackmodel exportarchive addendpoint nameserve fnmodel call inputsignature tf tensorspecshapenone 3 dtypetf float32 tf tensorspecshapenone 4 dtypetf float32 adding an endpoint using the inputsignature argument when the model has one input argument that is a list of 2 tensors e g a functional model with 2 inputs python model keras modelinputsx1 x2 outputsoutputs exportarchive exportarchive exportarchive trackmodel exportarchive addendpoint nameserve fnmodel call inputsignature tf tensorspecshapenone 3 dtypetf float32 tf tensorspecshapenone 4 dtypetf float32 this also works with dictionary inputs python model keras modelinputsx1 x1 x2 x2 outputsoutputs exportarchive exportarchive exportarchive trackmodel exportarchive addendpoint nameserve fnmodel call inputsignature x1 tf tensorspecshapenone 3 dtypetf float32 x2 tf tensorspecshapenone 4 dtypetf float32 adding an endpoint that is a tf function python tf function def servingfnx return modelx the function must be traced i e it must be called at least once servingfntf random normalshape2 3 exportarchive exportarchive exportarchive trackmodel exportarchive addendpointnameserve fnservingfn register a set of variables to be retrieved after reloading arguments name the string name for the collection variables a tuplelistset of tf variable instances example python exportarchive exportarchive exportarchive trackmodel register an endpoint exportarchive addendpoint nameserve fnmodel call inputsignaturetf tensorspecshapenone 3 dtypetf float32 save a variable collection exportarchive addvariablecollection nameoptimizervariables variablesmodel optimizer variables exportarchive writeoutpathtolocation reload the object revivedobject tf savedmodel loadpathtolocation retrieve the variables optimizervariables revivedobject optimizervariables ensure that all variables added are either tf variables or variables keras 3 with the tf or jax backends write the corresponding savedmodel to disk arguments filepath str or pathlib path object path where to save the artifact options tf savedmodel saveoptions object that specifies savedmodel saving options note on tfserving all endpoints registered via addendpoint are made visible for tfserving in the savedmodel artifact in addition the first endpoint registered is made visible under the alias servingdefault unless an endpoint with the name servingdefault was already registered manually since tfserving requires this endpoint to be set add servingdefault signature key for tfserving print out available endpoints workaround for some savedmodel quirks if endpoint in self endpointsignatures return getattrself tftrackable endpoint else traces getattrself tftrackable endpoint trackablechildren savedmodel return listtraces values0 def getvariablesusedbyendpointsself fns self getconcretefnname for name in self endpointnames return listvariablesusedbyfnsfns def filterandtrackresourcesself start by extracting variables from endpoints next track lookup tables hopefully one day this will be automated at the tf function level reload a keras modellayer that was saved via savedmodel exportarchive arguments filepath str or pathlib path object the path to the savedmodel callendpoint name of the endpoint to use as the call method of the reloaded layer if the savedmodel was created via model export then the default endpoint name is serve in other cases it may be named servingdefault example python model exportpathtoartifact reloadedlayer tfsmlayerpathtoartifact outputs reloadedlayerinputs the reloaded object can be used like a regular keras layer and supports trainingfinetuning of its trainable weights note that the reloaded object retains none of the internal structure or custom methods of the original object it s a brand new layer created around the saved function limitations only call endpoints with a single inputs tensor argument which may optionally be a dicttuplelist of tensors are supported for endpoints with multiple separate input tensor arguments consider subclassing tfsmlayer and implementing a call method with a custom signature if you need trainingtime behavior to differ from inferencetime behavior i e if you need the reloaded object to support a trainingtrue argument in call make sure that the trainingtime call function is saved as a standalone endpoint in the artifact and provide its name to the tfsmlayer via the calltrainingendpoint argument initialize an empty layer then addweight etc as needed resolve the call function case 1 it s set as an attribute case 2 it s listed in the signatures field resolving the training function add trainable and nontrainable weights from the callendpointfn tracks an existing weight self trackvariableweight def callself inputs trainingfalse kwargs if training if self calltrainingendpoint return self calltrainingendpointfninputs kwargs return self callendpointfninputs kwargs def getconfigself baseconfig super getconfig config note this is not intended to be portable filepath self filepath callendpoint self callendpoint calltrainingendpoint self calltrainingendpoint return baseconfig config def maketensorspecx return tf tensorspecx shape dtypex dtype namex name def printsignaturefn name concretefn fn listallconcretefunctions0 pprintedsignature concretefn prettyprintedsignatureverbosetrue lines pprintedsignature splitn lines f endpoint name lines1 endpoint n joinlines return endpoint def listvariablesusedbyfnsfns trainablevariables nontrainablevariables trainablevariablesids set nontrainablevariablesids set for fn in fns if hasattrfn concretefunctions concretefunctions fn concretefunctions elif hasattrfn getconcretefunction concretefunctions fn getconcretefunction else concretefunctions fn for concretefn in concretefunctions for v in concretefn trainablevariables if idv not in trainablevariablesids trainablevariables appendv trainablevariablesids addidv for v in concretefn variables if idv not in trainablevariablesids and idv not in nontrainablevariablesids nontrainablevariables appendv nontrainablevariablesids addidv return trainablevariables nontrainablevariables exportarchive is used to write savedmodel artifacts e g for inference if you have a keras model or layer that you want to export as savedmodel for serving e g via tensorflow serving you can use exportarchive to configure the different serving endpoints you need to make available as well as their signatures simply instantiate an exportarchive use track to register the layer s or model s to be used then use the add_endpoint method to register a new serving endpoint when done use the write_out method to save the artifact the resulting artifact is a savedmodel and can be reloaded via tf saved_model load examples here s how to export a model for inference python export_archive exportarchive export_archive track model export_archive add_endpoint name serve fn model call input_signature tf tensorspec shape none 3 dtype tf float32 export_archive write_out path to location elsewhere we can reload the artifact and serve it the endpoint we added is available as a method serving_model tf saved_model load path to location outputs serving_model serve inputs here s how to export a model with one endpoint for inference and one endpoint for a training mode forward pass e g with dropout on python export_archive exportarchive export_archive track model export_archive add_endpoint name call_inference fn lambda x model call x training false input_signature tf tensorspec shape none 3 dtype tf float32 export_archive add_endpoint name call_training fn lambda x model call x training true input_signature tf tensorspec shape none 3 dtype tf float32 export_archive write_out path to location note on resource tracking exportarchive is able to automatically track all tf variables used by its endpoints so most of the time calling track model is not strictly required however if your model uses lookup layers such as integerlookup stringlookup or textvectorization it will need to be tracked explicitly via track model explicit tracking is also required if you need to be able to access the properties variables trainable_variables or non_trainable_variables on the revived archive todo nkovela make jax version checking programatic track the variables and other assets of a layer or model layers in _tracked are not part of the trackables that get saved because we re creating the attribute in a no_automatic_dependency_tracking scope variables in the lists below are actually part of the trackables that get saved because the lists are created in __init__ register a new serving endpoint arguments name str name of the endpoint fn a function it should only leverage resources e g tf variable objects or tf lookup statichashtable objects that are available on the models layers tracked by the exportarchive you can call track model to track a new model the shape and dtype of the inputs to the function must be known for that purpose you can either 1 make sure that fn is a tf function that has been called at least once or 2 provide an input_signature argument that specifies the shape and dtype of the inputs see below input_signature used to specify the shape and dtype of the inputs to fn list of tf tensorspec objects one per positional input argument of fn nested arguments are allowed see below for an example showing a functional model with 2 input arguments example adding an endpoint using the input_signature argument when the model has a single input argument python export_archive exportarchive export_archive track model export_archive add_endpoint name serve fn model call input_signature tf tensorspec shape none 3 dtype tf float32 adding an endpoint using the input_signature argument when the model has two positional input arguments python export_archive exportarchive export_archive track model export_archive add_endpoint name serve fn model call input_signature tf tensorspec shape none 3 dtype tf float32 tf tensorspec shape none 4 dtype tf float32 adding an endpoint using the input_signature argument when the model has one input argument that is a list of 2 tensors e g a functional model with 2 inputs python model keras model inputs x1 x2 outputs outputs export_archive exportarchive export_archive track model export_archive add_endpoint name serve fn model call input_signature tf tensorspec shape none 3 dtype tf float32 tf tensorspec shape none 4 dtype tf float32 this also works with dictionary inputs python model keras model inputs x1 x1 x2 x2 outputs outputs export_archive exportarchive export_archive track model export_archive add_endpoint name serve fn model call input_signature x1 tf tensorspec shape none 3 dtype tf float32 x2 tf tensorspec shape none 4 dtype tf float32 adding an endpoint that is a tf function python tf function def serving_fn x return model x the function must be traced i e it must be called at least once serving_fn tf random normal shape 2 3 export_archive exportarchive export_archive track model export_archive add_endpoint name serve fn serving_fn jax backend register a set of variables to be retrieved after reloading arguments name the string name for the collection variables a tuple list set of tf variable instances example python export_archive exportarchive export_archive track model register an endpoint export_archive add_endpoint name serve fn model call input_signature tf tensorspec shape none 3 dtype tf float32 save a variable collection export_archive add_variable_collection name optimizer_variables variables model optimizer variables export_archive write_out path to location reload the object revived_object tf saved_model load path to location retrieve the variables optimizer_variables revived_object optimizer_variables ensure that all variables added are either tf variables or variables keras 3 with the tf or jax backends write the corresponding savedmodel to disk arguments filepath str or pathlib path object path where to save the artifact options tf saved_model saveoptions object that specifies savedmodel saving options note on tf serving all endpoints registered via add_endpoint are made visible for tf serving in the savedmodel artifact in addition the first endpoint registered is made visible under the alias serving_default unless an endpoint with the name serving_default was already registered manually since tf serving requires this endpoint to be set add serving_default signature key for tfserving print out available endpoints workaround for some savedmodel quirks track resources used by endpoints referenced in track calls start by extracting variables from endpoints next track lookup tables hopefully one day this will be automated at the tf function level reload a keras model layer that was saved via savedmodel exportarchive arguments filepath str or pathlib path object the path to the savedmodel call_endpoint name of the endpoint to use as the call method of the reloaded layer if the savedmodel was created via model export then the default endpoint name is serve in other cases it may be named serving_default example python model export path to artifact reloaded_layer tfsmlayer path to artifact outputs reloaded_layer inputs the reloaded object can be used like a regular keras layer and supports training fine tuning of its trainable weights note that the reloaded object retains none of the internal structure or custom methods of the original object it s a brand new layer created around the saved function limitations only call endpoints with a single inputs tensor argument which may optionally be a dict tuple list of tensors are supported for endpoints with multiple separate input tensor arguments consider subclassing tfsmlayer and implementing a call method with a custom signature if you need training time behavior to differ from inference time behavior i e if you need the reloaded object to support a training true argument in __call__ make sure that the training time call function is saved as a standalone endpoint in the artifact and provide its name to the tfsmlayer via the call_training_endpoint argument initialize an empty layer then add_weight etc as needed resolve the call function case 1 it s set as an attribute case 2 it s listed in the signatures field resolving the training function add trainable and non trainable weights from the call_endpoint_fn tracks an existing weight note this is not intended to be portable
from keras import backend
from keras.api_export import keras_export
from keras.layers import Layer
from keras.models import Functional
from keras.models import Sequential
from keras.utils import io_utils
from keras.utils.module_utils import tensorflow as tf


@keras_export("keras.export.ExportArchive")
class ExportArchive:
    

    def __init__(self):
        self._endpoint_names = []
        self._endpoint_signatures = {}
        self.tensorflow_version = tf.__version__

        self._tf_trackable = tf.__internal__.tracking.AutoTrackable()
        self._tf_trackable.variables = []
        self._tf_trackable.trainable_variables = []
        self._tf_trackable.non_trainable_variables = []

        if backend.backend() not in ("tensorflow", "jax"):
            raise NotImplementedError(
                "The export API is only compatible with JAX and TF backends."
            )

        
        if backend.backend() == "jax":
            from jax import __version__ as jax_v

            if jax_v > "0.4.15":
                raise ValueError(
                    "The export API is only compatible with JAX version 0.4.15 "
                    f"and prior. Your JAX version: {jax_v}"
                )

    @property
    def variables(self):
        return self._tf_trackable.variables

    @property
    def trainable_variables(self):
        return self._tf_trackable.trainable_variables

    @property
    def non_trainable_variables(self):
        return self._tf_trackable.non_trainable_variables

    def track(self, resource):
        
        if backend.backend() == "tensorflow" and not isinstance(
            resource, tf.__internal__.tracking.Trackable
        ):
            raise ValueError(
                "Invalid resource type. Expected an instance of a "
                "TensorFlow `Trackable` (such as a Keras `Layer` or `Model`). "
                f"Received instead an object of type '{type(resource)}'. "
                f"Object received: {resource}"
            )
        if backend.backend() == "jax" and not isinstance(
            resource, backend.jax.layer.JaxLayer
        ):
            raise ValueError(
                "Invalid resource type. Expected an instance of a "
                "JAX-based Keras `Layer` or `Model`. "
                f"Received instead an object of type '{type(resource)}'. "
                f"Object received: {resource}"
            )
        if isinstance(resource, Layer):
            if not resource.built:
                raise ValueError(
                    "The layer provided has not yet been built. "
                    "It must be built before export."
                )

        
        
        
        if not hasattr(self, "_tracked"):
            self._tracked = []
        self._tracked.append(resource)

        if isinstance(resource, Layer):
            
            
            self._tf_trackable.variables += resource.variables
            self._tf_trackable.trainable_variables += (
                resource.trainable_variables
            )
            self._tf_trackable.non_trainable_variables += (
                resource.non_trainable_variables
            )

    def add_endpoint(self, name, fn, input_signature=None):
        
        if name in self._endpoint_names:
            raise ValueError(f"Endpoint name '{name}' is already taken.")

        if input_signature:
            if backend.backend() == "tensorflow":
                decorated_fn = tf.function(fn, input_signature=input_signature)
            else:  
                fn = self._convert_jax2tf_function(fn, input_signature)
                decorated_fn = tf.function(
                    fn, input_signature=input_signature, autograph=False
                )
            self._endpoint_signatures[name] = input_signature
        else:
            if isinstance(fn, tf.types.experimental.GenericFunction):
                if not fn._list_all_concrete_functions():
                    raise ValueError(
                        f"The provided tf.function '{fn}' "
                        "has never been called. "
                        "To specify the expected shape and dtype "
                        "of the function's arguments, "
                        "you must either provide a function that "
                        "has been called at least once, or alternatively pass "
                        "an `input_signature` argument in `add_endpoint()`."
                    )
                decorated_fn = fn
            else:
                raise ValueError(
                    "If the `fn` argument provided is not a `tf.function`, "
                    "you must provide an `input_signature` argument to "
                    "specify the shape and dtype of the function arguments. "
                    "Example:\n\n"
                    "export_archive.add_endpoint(\n"
                    "    name='call',\n"
                    "    fn=model.call,\n"
                    "    input_signature=[\n"
                    "        tf.TensorSpec(\n"
                    "            shape=(None, 224, 224, 3),\n"
                    "            dtype=tf.float32,\n"
                    "        )\n"
                    "    ],\n"
                    ")"
                )
        setattr(self._tf_trackable, name, decorated_fn)
        self._endpoint_names.append(name)

    def add_variable_collection(self, name, variables):
        
        if not isinstance(variables, (list, tuple, set)):
            raise ValueError(
                "Expected `variables` to be a list/tuple/set. "
                f"Received instead object of type '{type(variables)}'."
            )
        
        
        if not all(
            isinstance(v, (tf.Variable, backend.Variable)) for v in variables
        ):
            raise ValueError(
                "Expected all elements in `variables` to be "
                "`tf.Variable` instances. Found instead the following types: "
                f"{list(set(type(v) for v in variables))}"
            )
        if backend.backend() == "jax":
            variables = tf.nest.flatten(
                tf.nest.map_structure(tf.Variable, variables)
            )
        setattr(self._tf_trackable, name, list(variables))

    def write_out(self, filepath, options=None):
        
        if not self._endpoint_names:
            raise ValueError(
                "No endpoints have been set yet. Call add_endpoint()."
            )
        if backend.backend() == "tensorflow":
            self._filter_and_track_resources()

        signatures = {}
        for name in self._endpoint_names:
            signatures[name] = self._get_concrete_fn(name)
        
        if "serving_default" not in self._endpoint_names:
            signatures["serving_default"] = self._get_concrete_fn(
                self._endpoint_names[0]
            )

        tf.saved_model.save(
            self._tf_trackable,
            filepath,
            options=options,
            signatures=signatures,
        )

        
        endpoints = "\n\n".join(
            _print_signature(getattr(self._tf_trackable, name), name)
            for name in self._endpoint_names
        )
        io_utils.print_msg(
            f"Saved artifact at '{filepath}'. "
            "The following endpoints are available:\n\n"
            f"{endpoints}"
        )

    def _get_concrete_fn(self, endpoint):
        
        if endpoint in self._endpoint_signatures:
            return getattr(self._tf_trackable, endpoint)
        else:
            traces = getattr(self._tf_trackable, endpoint)._trackable_children(
                "saved_model"
            )
            return list(traces.values())[0]

    def _get_variables_used_by_endpoints(self):
        fns = [self._get_concrete_fn(name) for name in self._endpoint_names]
        return _list_variables_used_by_fns(fns)

    def _filter_and_track_resources(self):
        
        
        fns = [self._get_concrete_fn(name) for name in self._endpoint_names]
        tvs, ntvs = _list_variables_used_by_fns(fns)
        self._tf_trackable._all_variables = list(tvs + ntvs)

        
        
        self._tf_trackable._misc_assets = []
        from keras.layers import IntegerLookup
        from keras.layers import StringLookup
        from keras.layers import TextVectorization

        if hasattr(self, "_tracked"):
            for root in self._tracked:
                descendants = tf.train.TrackableView(root).descendants()
                for trackable in descendants:
                    if isinstance(
                        trackable,
                        (IntegerLookup, StringLookup, TextVectorization),
                    ):
                        self._tf_trackable._misc_assets.append(trackable)

    def _convert_jax2tf_function(self, fn, input_signature):
        from jax.experimental import jax2tf

        shapes = []
        for spec in input_signature:
            shapes.append(self._spec_to_poly_shape(spec))
        return jax2tf.convert(fn, polymorphic_shapes=shapes)

    def _spec_to_poly_shape(self, spec):
        if isinstance(spec, (dict, list)):
            return tf.nest.map_structure(self._spec_to_poly_shape, spec)
        spec_shape = spec.shape
        spec_shape = str(spec_shape).replace("None", "b")
        return spec_shape


def export_model(model, filepath):
    export_archive = ExportArchive()
    export_archive.track(model)
    if isinstance(model, (Functional, Sequential)):
        input_signature = tf.nest.map_structure(_make_tensor_spec, model.inputs)
        if isinstance(input_signature, list) and len(input_signature) > 1:
            input_signature = [input_signature]
        export_archive.add_endpoint("serve", model.__call__, input_signature)
    else:
        save_spec = _get_save_spec(model)
        if not save_spec or not model._called:
            raise ValueError(
                "The model provided has never called. "
                "It must be called at least once before export."
            )
        input_signature = [save_spec]
        export_archive.add_endpoint("serve", model.__call__, input_signature)
    export_archive.write_out(filepath)


def _get_save_spec(model):
    shapes_dict = getattr(model, "_build_shapes_dict", None)
    if not shapes_dict:
        return None

    if len(shapes_dict) == 1:
        return tf.TensorSpec(
            shape=list(shapes_dict.values())[0], dtype=model.input_dtype
        )

    specs = {}
    for key, value in shapes_dict.items():
        key = key.rstrip("_shape")
        specs[key] = tf.TensorSpec(shape=value, dtype=model.input_dtype)
    return specs


@keras_export("keras.layers.TFSMLayer")
class TFSMLayer(Layer):
    

    def __init__(
        self,
        filepath,
        call_endpoint="serve",
        call_training_endpoint=None,
        trainable=True,
        name=None,
        dtype=None,
    ):
        
        super().__init__(trainable=trainable, name=name, dtype=dtype)

        self._reloaded_obj = tf.saved_model.load(filepath)

        self.filepath = filepath
        self.call_endpoint = call_endpoint
        self.call_training_endpoint = call_training_endpoint

        
        if hasattr(self._reloaded_obj, call_endpoint):
            
            self.call_endpoint_fn = getattr(self._reloaded_obj, call_endpoint)
        elif call_endpoint in self._reloaded_obj.signatures:
            
            self.call_endpoint_fn = self._reloaded_obj.signatures[call_endpoint]
        else:
            raise ValueError(
                f"The endpoint '{call_endpoint}' "
                "is neither an attribute of the reloaded SavedModel, "
                "nor an entry in the `signatures` field of "
                "the reloaded SavedModel. Select another endpoint via "
                "the `call_endpoint` argument. Available endpoints for "
                "this SavedModel: "
                f"{list(self._reloaded_obj.signatures.keys())}"
            )

        
        if call_training_endpoint:
            if hasattr(self._reloaded_obj, call_training_endpoint):
                self.call_training_endpoint_fn = getattr(
                    self._reloaded_obj, call_training_endpoint
                )
            elif call_training_endpoint in self._reloaded_obj.signatures:
                self.call_training_endpoint_fn = self._reloaded_obj.signatures[
                    call_training_endpoint
                ]
            else:
                raise ValueError(
                    f"The endpoint '{call_training_endpoint}' "
                    "is neither an attribute of the reloaded SavedModel, "
                    "nor an entry in the `signatures` field of "
                    "the reloaded SavedModel. Available endpoints for "
                    "this SavedModel: "
                    f"{list(self._reloaded_obj.signatures.keys())}"
                )

        
        all_fns = [self.call_endpoint_fn]
        if call_training_endpoint:
            all_fns.append(self.call_training_endpoint_fn)
        tvs, ntvs = _list_variables_used_by_fns(all_fns)
        for v in tvs:
            self._add_existing_weight(v)
        for v in ntvs:
            self._add_existing_weight(v)
        self.built = True

    def _add_existing_weight(self, weight):
        
        self._track_variable(weight)

    def call(self, inputs, training=False, **kwargs):
        if training:
            if self.call_training_endpoint:
                return self.call_training_endpoint_fn(inputs, **kwargs)
        return self.call_endpoint_fn(inputs, **kwargs)

    def get_config(self):
        base_config = super().get_config()
        config = {
            
            "filepath": self.filepath,
            "call_endpoint": self.call_endpoint,
            "call_training_endpoint": self.call_training_endpoint,
        }
        return {**base_config, **config}


def _make_tensor_spec(x):
    return tf.TensorSpec(x.shape, dtype=x.dtype, name=x.name)


def _print_signature(fn, name):
    concrete_fn = fn._list_all_concrete_functions()[0]
    pprinted_signature = concrete_fn.pretty_printed_signature(verbose=True)
    lines = pprinted_signature.split("\n")
    lines = [f"* Endpoint '{name}'"] + lines[1:]
    endpoint = "\n".join(lines)
    return endpoint


def _list_variables_used_by_fns(fns):
    trainable_variables = []
    non_trainable_variables = []
    trainable_variables_ids = set()
    non_trainable_variables_ids = set()
    for fn in fns:
        if hasattr(fn, "concrete_functions"):
            concrete_functions = fn.concrete_functions
        elif hasattr(fn, "get_concrete_function"):
            concrete_functions = [fn.get_concrete_function()]
        else:
            concrete_functions = [fn]
        for concrete_fn in concrete_functions:
            for v in concrete_fn.trainable_variables:
                if id(v) not in trainable_variables_ids:
                    trainable_variables.append(v)
                    trainable_variables_ids.add(id(v))

            for v in concrete_fn.variables:
                if (
                    id(v) not in trainable_variables_ids
                    and id(v) not in non_trainable_variables_ids
                ):
                    non_trainable_variables.append(v)
                    non_trainable_variables_ids.add(id(v))
    return trainable_variables, non_trainable_variables