tests for inferenceonly modellayer exporting utilities import os import sys import numpy as np import pytest import tensorflow as tf from keras import backend from keras import layers from keras import models from keras import testing from keras import utils from keras export import exportlib from keras saving import savinglib def getmodel layerlist layers dense10 activationrelu layers batchnormalization layers dense1 activationsigmoid model models sequentiallayerlist return model pytest mark skipif backend backend not in tensorflow jax reasonexport only currently supports the tf and jax backends pytest mark skipif backend backend jax and sys modulesjax version 0 4 15 reasonthe export api is only compatible with jax version 0 4 15 class exportarchivetesttesting testcase def teststandardmodelexportself tempfilepath os path joinself gettempdir exportedmodel model getmodel refinput tf random normal3 10 refoutput modelrefinput exportlib exportmodelmodel tempfilepath revivedmodel tf savedmodel loadtempfilepath self assertallclose refoutput revivedmodel serverefinput atol1e6 def testlowlevelmodelexportself tempfilepath os path joinself gettempdir exportedmodel model getmodel refinput tf random normal3 10 refoutput modelrefinput test variable tracking exportarchive exportlib exportarchive exportarchive trackmodel self assertlenexportarchive variables 8 self assertlenexportarchive trainablevariables 6 self assertlenexportarchive nontrainablevariables 2 exportarchive exportlib exportarchive exportarchive trackmodel exportarchive addendpoint call model call inputsignature tf tensorspec shapenone 10 dtypetf float32 exportarchive writeouttempfilepath revivedmodel tf savedmodel loadtempfilepath self assertallclose refoutput revivedmodel callrefinput atol1e6 pytest mark skipif backend backend tensorflow reasonregistering a tf function endpoint is only in tf backend def testendpointregistrationtffunctionself tempfilepath os path joinself gettempdir exportedmodel model getmodel refinput tf random normal3 10 refoutput modelrefinput test variable tracking exportarchive exportlib exportarchive exportarchive trackmodel self assertlenexportarchive variables 8 self assertlenexportarchive trainablevariables 6 self assertlenexportarchive nontrainablevariables 2 tf function def myendpointx return modelx test registering an endpoint that is a tf function called myendpointrefinput trace fn exportarchive addendpoint call myendpoint exportarchive writeouttempfilepath revivedmodel tf savedmodel loadtempfilepath self assertfalsehasattrrevivedmodel tracked self assertallclose refoutput revivedmodel callrefinput atol1e6 self assertlenrevivedmodel variables 8 self assertlenrevivedmodel trainablevariables 6 self assertlenrevivedmodel nontrainablevariables 2 def testlayerexportself tempfilepath os path joinself gettempdir exportedlayer layer layers batchnormalization refinput tf random normal3 10 refoutput layerrefinput build layer important exportarchive exportlib exportarchive exportarchive tracklayer exportarchive addendpoint call layer call inputsignature tf tensorspec shapenone 10 dtypetf float32 exportarchive writeouttempfilepath revivedlayer tf savedmodel loadtempfilepath self assertallclose refoutput revivedlayer callrefinput atol1e6 def testmultiinputoutputfunctionalmodelself tempfilepath os path joinself gettempdir exportedmodel x1 layers input2 x2 layers input2 y1 layers dense3x1 y2 layers dense3x2 model models modelx1 x2 y1 y2 refinputs tf random normal3 2 tf random normal3 2 refoutputs modelrefinputs exportarchive exportlib exportarchive exportarchive trackmodel exportarchive addendpoint serve model call inputsignature tf tensorspec shapenone 2 dtypetf float32 tf tensorspec shapenone 2 dtypetf float32 exportarchive writeouttempfilepath revivedmodel tf savedmodel loadtempfilepath self assertallclose refoutputs0 revivedmodel serverefinputs0 atol1e6 self assertallclose refoutputs1 revivedmodel serverefinputs1 atol1e6 now test dict inputs model models modelx1 x1 x2 x2 y1 y2 refinputs x1 tf random normal3 2 x2 tf random normal3 2 refoutputs modelrefinputs exportarchive exportlib exportarchive exportarchive trackmodel exportarchive addendpoint serve model call inputsignature x1 tf tensorspec shapenone 2 dtypetf float32 x2 tf tensorspec shapenone 2 dtypetf float32 exportarchive writeouttempfilepath revivedmodel tf savedmodel loadtempfilepath self assertallclose refoutputs0 revivedmodel serverefinputs0 atol1e6 self assertallclose refoutputs1 revivedmodel serverefinputs1 atol1e6 def testmodelwithlookuptableself tf debugging disabletracebackfiltering tempfilepath os path joinself gettempdir exportedmodel textvectorization layers textvectorization textvectorization adaptone two three four five six model models sequential layers inputshape dtypestring textvectorization layers embedding10 32 layers dense1 refinput tf converttotensorone two three four refoutput modelrefinput exportlib exportmodelmodel tempfilepath revivedmodel tf savedmodel loadtempfilepath self assertallclose refoutput revivedmodel serverefinput atol1e6 def testtrackmultiplelayersself tempfilepath os path joinself gettempdir exportedmodel layer1 layers dense2 refinput1 tf random normal3 4 refoutput1 layer1refinput1 layer2 layers dense3 refinput2 tf random normal3 5 refoutput2 layer2refinput2 exportarchive exportlib exportarchive exportarchive addendpoint call1 layer1 call inputsignature tf tensorspec shapenone 4 dtypetf float32 exportarchive addendpoint call2 layer2 call inputsignature tf tensorspec shapenone 5 dtypetf float32 exportarchive writeouttempfilepath revivedlayer tf savedmodel loadtempfilepath self assertallclose refoutput1 revivedlayer call1refinput1 atol1e6 self assertallclose refoutput2 revivedlayer call2refinput2 atol1e6 def testnonstandardlayersignatureself tempfilepath os path joinself gettempdir exportedlayer layer layers multiheadattention2 2 x1 tf random normal3 2 2 x2 tf random normal3 2 2 refoutput layerx1 x2 build layer important exportarchive exportlib exportarchive exportarchive tracklayer exportarchive addendpoint call layer call inputsignature tf tensorspec shapenone 2 2 dtypetf float32 tf tensorspec shapenone 2 2 dtypetf float32 exportarchive writeouttempfilepath revivedlayer tf savedmodel loadtempfilepath self assertallclose refoutput revivedlayer callx1 x2 atol1e6 todonkovela remove test when argument name preservation workaround is created for jax backend pytest mark skipif backend backend tensorflow reasonjax2tf has issues with argument name preservation def testnonstandardlayersignaturewithkwargsself tempfilepath os path joinself gettempdir exportedlayer layer layers multiheadattention2 2 x1 tf random normal3 2 2 x2 tf random normal3 2 2 refoutput layerx1 x2 build layer important exportarchive exportlib exportarchive exportarchive tracklayer exportarchive addendpoint call layer call inputsignature tf tensorspec shapenone 2 2 dtypetf float32 tf tensorspec shapenone 2 2 dtypetf float32 exportarchive writeouttempfilepath revivedlayer tf savedmodel loadtempfilepath self assertallclose refoutput revivedlayer callqueryx1 valuex2 atol1e6 def testvariablecollectionself tempfilepath os path joinself gettempdir exportedmodel model models sequential layers input10 layers dense2 layers dense2 test variable tracking exportarchive exportlib exportarchive exportarchive trackmodel exportarchive addendpoint call model call inputsignature tf tensorspec shapenone 10 dtypetf float32 exportarchive addvariablecollection myvars model layers1 weights self assertlenexportarchive tftrackable myvars 2 exportarchive writeouttempfilepath revivedmodel tf savedmodel loadtempfilepath self assertlenrevivedmodel myvars 2 def testexportmodelerrorsself tempfilepath os path joinself gettempdir exportedmodel model has not been built model models sequentiallayers dense2 with self assertraisesregexvalueerror it must be built exportlib exportmodelmodel tempfilepath subclassed model has not been called class mymodelmodels model def initself kwargs super initkwargs self dense layers dense2 def buildself inputshape self dense buildinputshape self built true def callself x return self densex model mymodel model build2 3 with self assertraisesregexvalueerror it must be called exportlib exportmodelmodel tempfilepath def testexportarchiveerrorsself tempfilepath os path joinself gettempdir exportedmodel model models sequentiallayers dense2 modeltf random normal2 3 endpoint name reuse exportarchive exportlib exportarchive exportarchive trackmodel exportarchive addendpoint call model call inputsignature tf tensorspec shapenone 3 dtypetf float32 with self assertraisesregexvalueerror already taken exportarchive addendpoint call model call inputsignature tf tensorspec shapenone 3 dtypetf float32 write out with no endpoints exportarchive exportlib exportarchive exportarchive trackmodel with self assertraisesregexvalueerror no endpoints have been set exportarchive writeouttempfilepath invalid object type with self assertraisesregexvalueerror invalid resource type exportarchive exportlib exportarchive exportarchive trackmodel set endpoint with no input signature exportarchive exportlib exportarchive exportarchive trackmodel with self assertraisesregex valueerror you must provide an inputsignature exportarchive addendpoint call model call set endpoint that has never been called exportarchive exportlib exportarchive exportarchive trackmodel tf function def myendpointx return modelx exportarchive exportlib exportarchive exportarchive trackmodel with self assertraisesregex valueerror you must either provide a function exportarchive addendpoint call myendpoint def testsubclassedmodelexportself class custommodelxmodels model def initself args kwargs super initargs kwargs self dense1 layers dense1 self dense2 layers dense1 def callself inputs out self dense1inputs return self dense2out tempfilepath os path joinself gettempdir exportedmodel x np random random100 32 model custommodelx model compile optimizeradam lossmse refoutput modelx model exporttempfilepath revivedmodel tf savedmodel loadtempfilepath self assertallcloserefoutput revivedmodel servex atol1e6 def testexportnoassetsself tempfilepath os path joinself gettempdir exportedmodel case where there are legitimately no assets model models sequentiallayers flatten modeltf random normal2 3 exportarchive exportlib exportarchive exportarchive addendpoint call model call inputsignature tf tensorspec shapenone 3 dtypetf float32 exportarchive writeouttempfilepath def testmodelexportmethodself tempfilepath os path joinself gettempdir exportedmodel model getmodel refinput tf random normal3 10 refoutput modelrefinput model exporttempfilepath revivedmodel tf savedmodel loadtempfilepath self assertallclose refoutput revivedmodel serverefinput atol1e6 pytest mark skipif backend backend jax or sys modulesjax version 0 4 15 reasonthis test is for invalid jax versions i e versions 0 4 15 class versiontesttesting testcase def testinvalidjaxversionself with self assertraisesregex valueerror only compatible with jax version exportlib exportarchive pytest mark skipif backend backend tensorflow reasontfsm layer reloading is only for the tf backend class testtfsmlayertesting testcase def testreloadingexportarchiveself tempfilepath os path joinself gettempdir exportedmodel model getmodel refinput tf random normal3 10 refoutput modelrefinput exportlib exportmodelmodel tempfilepath reloadedlayer exportlib tfsmlayertempfilepath self assertallclosereloadedlayerrefinput refoutput atol1e7 self assertlenreloadedlayer weights lenmodel weights self assertlen reloadedlayer trainableweights lenmodel trainableweights self assertlen reloadedlayer nontrainableweights lenmodel nontrainableweights todonkovela expand test coveragedebug finetuning and nontrainable use cases here def testreloadingdefaultsavedmodelself tempfilepath os path joinself gettempdir exportedmodel model getmodel refinput tf random normal3 10 refoutput modelrefinput tf savedmodel savemodel tempfilepath reloadedlayer exportlib tfsmlayer tempfilepath callendpointservingdefault the output is a dict due to the nature of savedmodel saving newoutput reloadedlayerrefinput self assertallclose newoutputlistnewoutput keys0 refoutput atol1e7 self assertlenreloadedlayer weights lenmodel weights self assertlen reloadedlayer trainableweights lenmodel trainableweights self assertlen reloadedlayer nontrainableweights lenmodel nontrainableweights def testcalltrainingself tempfilepath os path joinself gettempdir exportedmodel utils setrandomseed1337 model models sequential layers input10 layers dense10 layers dropout0 99999 exportarchive exportlib exportarchive exportarchive trackmodel exportarchive addendpoint namecallinference fnlambda x modelx trainingfalse inputsignaturetf tensorspecshapenone 10 dtypetf float32 exportarchive addendpoint namecalltraining fnlambda x modelx trainingtrue inputsignaturetf tensorspecshapenone 10 dtypetf float32 exportarchive writeouttempfilepath reloadedlayer exportlib tfsmlayer tempfilepath callendpointcallinference calltrainingendpointcalltraining inferenceoutput reloadedlayer tf random normal1 10 trainingfalse trainingoutput reloadedlayer tf random normal1 10 trainingtrue self assertallclosenp meantrainingoutput 0 0 atol1e7 self assertnotallclosenp meaninferenceoutput 0 0 atol1e7 def testserializationself tempfilepath os path joinself gettempdir exportedmodel model getmodel refinput tf random normal3 10 refoutput modelrefinput exportlib exportmodelmodel tempfilepath reloadedlayer exportlib tfsmlayertempfilepath test reinstantiation from config config reloadedlayer getconfig rereloadedlayer exportlib tfsmlayer fromconfigconfig self assertallcloserereloadedlayerrefinput refoutput atol1e7 test whole model saving with reloaded layer inside model models sequentialreloadedlayer tempmodelfilepath os path joinself gettempdir m keras model savetempmodelfilepath saveformatkerasv3 reloadedmodel savinglib loadmodel tempmodelfilepath customobjectstfsmlayer exportlib tfsmlayer self assertallclosereloadedmodelrefinput refoutput atol1e7 def testerrorsself test missing call endpoint tempfilepath os path joinself gettempdir exportedmodel model models sequentiallayers input2 layers dense3 exportlib exportmodelmodel tempfilepath with self assertraisesregexvalueerror the endpoint wrong exportlib tfsmlayertempfilepath callendpointwrong test missing call training endpoint with self assertraisesregexvalueerror the endpoint wrong exportlib tfsmlayer tempfilepath callendpointserve calltrainingendpointwrong test variable tracking test variable tracking test registering an endpoint that is a tf function called trace fn build layer important now test dict inputs def test_model_with_lookup_table self tf debugging disable_traceback_filtering temp_filepath os path join self get_temp_dir exported_model text_vectorization layers textvectorization text_vectorization adapt one two three four five six model models sequential layers input shape dtype string text_vectorization layers embedding 10 32 layers dense 1 ref_input tf convert_to_tensor one two three four ref_output model ref_input export_lib export_model model temp_filepath revived_model tf saved_model load temp_filepath self assertallclose ref_output revived_model serve ref_input atol 1e 6 build layer important todo nkovela remove test when argument name preservation workaround is created for jax backend build layer important test variable tracking model has not been built subclassed model has not been called endpoint name reuse write out with no endpoints invalid object type set endpoint with no input signature set endpoint that has never been called case where there are legitimately no assets todo nkovela expand test coverage debug fine tuning and non trainable use cases here the output is a dict due to the nature of savedmodel saving test reinstantiation from config test whole model saving with reloaded layer inside test missing call endpoint test missing call training endpoint
import os
import sys

import numpy as np
import pytest
import tensorflow as tf

from keras import backend
from keras import layers
from keras import models
from keras import testing
from keras import utils
from keras.export import export_lib
from keras.saving import saving_lib


def get_model():
    layer_list = [
        layers.Dense(10, activation="relu"),
        layers.BatchNormalization(),
        layers.Dense(1, activation="sigmoid"),
    ]
    model = models.Sequential(layer_list)
    return model


@pytest.mark.skipif(
    backend.backend() not in ("tensorflow", "jax"),
    reason="Export only currently supports the TF and JAX backends.",
)
@pytest.mark.skipif(
    backend.backend() == "jax" and sys.modules["jax"].__version__ > "0.4.15",
    reason="The export API is only compatible with JAX version <= 0.4.15.",
)
class ExportArchiveTest(testing.TestCase):
    def test_standard_model_export(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_lib.export_model(model, temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(
            ref_output, revived_model.serve(ref_input), atol=1e-6
        )

    def test_low_level_model_export(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        self.assertLen(export_archive.variables, 8)
        self.assertLen(export_archive.trainable_variables, 6)
        self.assertLen(export_archive.non_trainable_variables, 2)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.call,
            input_signature=[
                tf.TensorSpec(
                    shape=(None, 10),
                    dtype=tf.float32,
                )
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(
            ref_output, revived_model.call(ref_input), atol=1e-6
        )

    @pytest.mark.skipif(
        backend.backend() != "tensorflow",
        reason="Registering a tf.function endpoint is only in TF backend.",
    )
    def test_endpoint_registration_tf_function(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        self.assertLen(export_archive.variables, 8)
        self.assertLen(export_archive.trainable_variables, 6)
        self.assertLen(export_archive.non_trainable_variables, 2)

        @tf.function()
        def my_endpoint(x):
            return model(x)

        
        my_endpoint(ref_input)  

        export_archive.add_endpoint(
            "call",
            my_endpoint,
        )
        export_archive.write_out(temp_filepath)

        revived_model = tf.saved_model.load(temp_filepath)
        self.assertFalse(hasattr(revived_model, "_tracked"))
        self.assertAllClose(
            ref_output, revived_model.call(ref_input), atol=1e-6
        )
        self.assertLen(revived_model.variables, 8)
        self.assertLen(revived_model.trainable_variables, 6)
        self.assertLen(revived_model.non_trainable_variables, 2)

    def test_layer_export(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_layer")

        layer = layers.BatchNormalization()
        ref_input = tf.random.normal((3, 10))
        ref_output = layer(ref_input)  

        export_archive = export_lib.ExportArchive()
        export_archive.track(layer)
        export_archive.add_endpoint(
            "call",
            layer.call,
            input_signature=[
                tf.TensorSpec(
                    shape=(None, 10),
                    dtype=tf.float32,
                )
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.saved_model.load(temp_filepath)
        self.assertAllClose(
            ref_output, revived_layer.call(ref_input), atol=1e-6
        )

    def test_multi_input_output_functional_model(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        x1 = layers.Input((2,))
        x2 = layers.Input((2,))
        y1 = layers.Dense(3)(x1)
        y2 = layers.Dense(3)(x2)
        model = models.Model([x1, x2], [y1, y2])

        ref_inputs = [tf.random.normal((3, 2)), tf.random.normal((3, 2))]
        ref_outputs = model(ref_inputs)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "serve",
            model.call,
            input_signature=[
                [
                    tf.TensorSpec(
                        shape=(None, 2),
                        dtype=tf.float32,
                    ),
                    tf.TensorSpec(
                        shape=(None, 2),
                        dtype=tf.float32,
                    ),
                ]
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(
            ref_outputs[0],
            revived_model.serve(ref_inputs)[0],
            atol=1e-6,
        )
        self.assertAllClose(
            ref_outputs[1],
            revived_model.serve(ref_inputs)[1],
            atol=1e-6,
        )

        
        model = models.Model({"x1": x1, "x2": x2}, [y1, y2])

        ref_inputs = {
            "x1": tf.random.normal((3, 2)),
            "x2": tf.random.normal((3, 2)),
        }
        ref_outputs = model(ref_inputs)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "serve",
            model.call,
            input_signature=[
                {
                    "x1": tf.TensorSpec(
                        shape=(None, 2),
                        dtype=tf.float32,
                    ),
                    "x2": tf.TensorSpec(
                        shape=(None, 2),
                        dtype=tf.float32,
                    ),
                }
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(
            ref_outputs[0],
            revived_model.serve(ref_inputs)[0],
            atol=1e-6,
        )
        self.assertAllClose(
            ref_outputs[1],
            revived_model.serve(ref_inputs)[1],
            atol=1e-6,
        )

    
    
    
    
    
    
    
    
    
    
    
    
    
    
    

    
    
    
    
    

    def test_track_multiple_layers(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        layer_1 = layers.Dense(2)
        ref_input_1 = tf.random.normal((3, 4))
        ref_output_1 = layer_1(ref_input_1)
        layer_2 = layers.Dense(3)
        ref_input_2 = tf.random.normal((3, 5))
        ref_output_2 = layer_2(ref_input_2)

        export_archive = export_lib.ExportArchive()
        export_archive.add_endpoint(
            "call_1",
            layer_1.call,
            input_signature=[
                tf.TensorSpec(
                    shape=(None, 4),
                    dtype=tf.float32,
                ),
            ],
        )
        export_archive.add_endpoint(
            "call_2",
            layer_2.call,
            input_signature=[
                tf.TensorSpec(
                    shape=(None, 5),
                    dtype=tf.float32,
                ),
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.saved_model.load(temp_filepath)
        self.assertAllClose(
            ref_output_1,
            revived_layer.call_1(ref_input_1),
            atol=1e-6,
        )
        self.assertAllClose(
            ref_output_2,
            revived_layer.call_2(ref_input_2),
            atol=1e-6,
        )

    def test_non_standard_layer_signature(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_layer")

        layer = layers.MultiHeadAttention(2, 2)
        x1 = tf.random.normal((3, 2, 2))
        x2 = tf.random.normal((3, 2, 2))
        ref_output = layer(x1, x2)  
        export_archive = export_lib.ExportArchive()
        export_archive.track(layer)
        export_archive.add_endpoint(
            "call",
            layer.call,
            input_signature=[
                tf.TensorSpec(
                    shape=(None, 2, 2),
                    dtype=tf.float32,
                ),
                tf.TensorSpec(
                    shape=(None, 2, 2),
                    dtype=tf.float32,
                ),
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.saved_model.load(temp_filepath)
        self.assertAllClose(
            ref_output,
            revived_layer.call(x1, x2),
            atol=1e-6,
        )

    
    
    @pytest.mark.skipif(
        backend.backend() != "tensorflow",
        reason="JAX2TF has issues with argument name preservation.",
    )
    def test_non_standard_layer_signature_with_kwargs(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_layer")

        layer = layers.MultiHeadAttention(2, 2)
        x1 = tf.random.normal((3, 2, 2))
        x2 = tf.random.normal((3, 2, 2))
        ref_output = layer(x1, x2)  
        export_archive = export_lib.ExportArchive()
        export_archive.track(layer)
        export_archive.add_endpoint(
            "call",
            layer.call,
            input_signature=[
                tf.TensorSpec(
                    shape=(None, 2, 2),
                    dtype=tf.float32,
                ),
                tf.TensorSpec(
                    shape=(None, 2, 2),
                    dtype=tf.float32,
                ),
            ],
        )
        export_archive.write_out(temp_filepath)
        revived_layer = tf.saved_model.load(temp_filepath)
        self.assertAllClose(
            ref_output,
            revived_layer.call(query=x1, value=x2),
            atol=1e-6,
        )

    def test_variable_collection(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        model = models.Sequential(
            [
                layers.Input((10,)),
                layers.Dense(2),
                layers.Dense(2),
            ]
        )

        
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.call,
            input_signature=[
                tf.TensorSpec(
                    shape=(None, 10),
                    dtype=tf.float32,
                )
            ],
        )
        export_archive.add_variable_collection(
            "my_vars", model.layers[1].weights
        )

        self.assertLen(export_archive._tf_trackable.my_vars, 2)
        export_archive.write_out(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertLen(revived_model.my_vars, 2)

    def test_export_model_errors(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        
        model = models.Sequential([layers.Dense(2)])
        with self.assertRaisesRegex(ValueError, "It must be built"):
            export_lib.export_model(model, temp_filepath)

        
        class MyModel(models.Model):
            def __init__(self, **kwargs):
                super().__init__(**kwargs)
                self.dense = layers.Dense(2)

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

            def call(self, x):
                return self.dense(x)

        model = MyModel()
        model.build((2, 3))
        with self.assertRaisesRegex(ValueError, "It must be called"):
            export_lib.export_model(model, temp_filepath)

    def test_export_archive_errors(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = models.Sequential([layers.Dense(2)])
        model(tf.random.normal((2, 3)))

        
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            "call",
            model.call,
            input_signature=[
                tf.TensorSpec(
                    shape=(None, 3),
                    dtype=tf.float32,
                )
            ],
        )
        with self.assertRaisesRegex(ValueError, "already taken"):
            export_archive.add_endpoint(
                "call",
                model.call,
                input_signature=[
                    tf.TensorSpec(
                        shape=(None, 3),
                        dtype=tf.float32,
                    )
                ],
            )

        
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        with self.assertRaisesRegex(ValueError, "No endpoints have been set"):
            export_archive.write_out(temp_filepath)

        
        with self.assertRaisesRegex(ValueError, "Invalid resource type"):
            export_archive = export_lib.ExportArchive()
            export_archive.track("model")

        
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        with self.assertRaisesRegex(
            ValueError, "you must provide an `input_signature`"
        ):
            export_archive.add_endpoint(
                "call",
                model.call,
            )

        
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)

        @tf.function()
        def my_endpoint(x):
            return model(x)

        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        with self.assertRaisesRegex(
            ValueError, "you must either provide a function"
        ):
            export_archive.add_endpoint(
                "call",
                my_endpoint,
            )

    def test_subclassed_model_export(self):
        class CustomModelX(models.Model):
            def __init__(self, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.dense1 = layers.Dense(1)
                self.dense2 = layers.Dense(1)

            def call(self, inputs):
                out = self.dense1(inputs)
                return self.dense2(out)

        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        x = np.random.random((100, 32))
        model = CustomModelX()
        model.compile(
            optimizer="adam",
            loss="mse",
        )
        ref_output = model(x)
        model.export(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.serve(x), atol=1e-6)

    def test_export_no_assets(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")

        
        model = models.Sequential([layers.Flatten()])
        model(tf.random.normal((2, 3)))
        export_archive = export_lib.ExportArchive()
        export_archive.add_endpoint(
            "call",
            model.call,
            input_signature=[
                tf.TensorSpec(
                    shape=(None, 3),
                    dtype=tf.float32,
                )
            ],
        )
        export_archive.write_out(temp_filepath)

    def test_model_export_method(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        model.export(temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(
            ref_output, revived_model.serve(ref_input), atol=1e-6
        )


@pytest.mark.skipif(
    backend.backend() != "jax" or sys.modules["jax"].__version__ <= "0.4.15",
    reason="This test is for invalid JAX versions, i.e. versions > 0.4.15.",
)
class VersionTest(testing.TestCase):
    def test_invalid_jax_version(self):
        with self.assertRaisesRegex(
            ValueError, "only compatible with JAX version"
        ):
            _ = export_lib.ExportArchive()


@pytest.mark.skipif(
    backend.backend() != "tensorflow",
    reason="TFSM Layer reloading is only for the TF backend.",
)
class TestTFSMLayer(testing.TestCase):
    def test_reloading_export_archive(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_lib.export_model(model, temp_filepath)
        reloaded_layer = export_lib.TFSMLayer(temp_filepath)
        self.assertAllClose(reloaded_layer(ref_input), ref_output, atol=1e-7)
        self.assertLen(reloaded_layer.weights, len(model.weights))
        self.assertLen(
            reloaded_layer.trainable_weights, len(model.trainable_weights)
        )
        self.assertLen(
            reloaded_layer.non_trainable_weights,
            len(model.non_trainable_weights),
        )

        
        

    def test_reloading_default_saved_model(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        tf.saved_model.save(model, temp_filepath)
        reloaded_layer = export_lib.TFSMLayer(
            temp_filepath, call_endpoint="serving_default"
        )
        
        new_output = reloaded_layer(ref_input)
        self.assertAllClose(
            new_output[list(new_output.keys())[0]],
            ref_output,
            atol=1e-7,
        )
        self.assertLen(reloaded_layer.weights, len(model.weights))
        self.assertLen(
            reloaded_layer.trainable_weights, len(model.trainable_weights)
        )
        self.assertLen(
            reloaded_layer.non_trainable_weights,
            len(model.non_trainable_weights),
        )

    def test_call_training(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        utils.set_random_seed(1337)
        model = models.Sequential(
            [
                layers.Input((10,)),
                layers.Dense(10),
                layers.Dropout(0.99999),
            ]
        )
        export_archive = export_lib.ExportArchive()
        export_archive.track(model)
        export_archive.add_endpoint(
            name="call_inference",
            fn=lambda x: model(x, training=False),
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
        )
        export_archive.add_endpoint(
            name="call_training",
            fn=lambda x: model(x, training=True),
            input_signature=[tf.TensorSpec(shape=(None, 10), dtype=tf.float32)],
        )
        export_archive.write_out(temp_filepath)
        reloaded_layer = export_lib.TFSMLayer(
            temp_filepath,
            call_endpoint="call_inference",
            call_training_endpoint="call_training",
        )
        inference_output = reloaded_layer(
            tf.random.normal((1, 10)), training=False
        )
        training_output = reloaded_layer(
            tf.random.normal((1, 10)), training=True
        )
        self.assertAllClose(np.mean(training_output), 0.0, atol=1e-7)
        self.assertNotAllClose(np.mean(inference_output), 0.0, atol=1e-7)

    def test_serialization(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model()
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_lib.export_model(model, temp_filepath)
        reloaded_layer = export_lib.TFSMLayer(temp_filepath)

        
        config = reloaded_layer.get_config()
        rereloaded_layer = export_lib.TFSMLayer.from_config(config)
        self.assertAllClose(rereloaded_layer(ref_input), ref_output, atol=1e-7)

        
        model = models.Sequential([reloaded_layer])
        temp_model_filepath = os.path.join(self.get_temp_dir(), "m.keras")
        model.save(temp_model_filepath, save_format="keras_v3")
        reloaded_model = saving_lib.load_model(
            temp_model_filepath,
            custom_objects={"TFSMLayer": export_lib.TFSMLayer},
        )
        self.assertAllClose(reloaded_model(ref_input), ref_output, atol=1e-7)

    def test_errors(self):
        
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = models.Sequential([layers.Input((2,)), layers.Dense(3)])
        export_lib.export_model(model, temp_filepath)
        with self.assertRaisesRegex(ValueError, "The endpoint 'wrong'"):
            export_lib.TFSMLayer(temp_filepath, call_endpoint="wrong")

        
        with self.assertRaisesRegex(ValueError, "The endpoint 'wrong'"):
            export_lib.TFSMLayer(
                temp_filepath,
                call_endpoint="serve",
                call_training_endpoint="wrong",
            )