Don't save traces for saved model

We check compilation elsewhere, adding this option significantly speeds
up saved model testing. Roughly 3x for the tf format.

Author: mattdangerw

keras_nlp/layers/f_net_encoder_test.py

@@ -133,7 +133,7 @@ def test_saved_model(self, save_format, filename):
         data = tf.random.uniform(shape=[2, 4, 6])
         model(data)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        model.save(path, save_traces=False, save_format=save_format)
         loaded_model = keras.models.load_model(path)
 
         model_output = model(data)

keras_nlp/layers/multi_segment_packer_test.py

@@ -178,7 +178,7 @@ def test_saved_model(self, save_format, filename):
         outputs = packer(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        model.save(path, save_traces=False, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model((seq1, seq2)),

keras_nlp/layers/position_embedding_test.py

@@ -309,7 +309,7 @@ def test_saved_model(self, save_format, filename):
         model(data)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        model.save(path, save_traces=False, save_format=save_format)
         loaded_model = keras.models.load_model(path)
 
         model_output = model.predict(data)

keras_nlp/layers/token_and_position_embedding_test.py

@@ -148,7 +148,7 @@ def test_saved_model(self, save_format, filename):
         model(data)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        model.save(path, save_traces=False, save_format=save_format)
         loaded_model = keras.models.load_model(path)
 
         model_output = model.predict(data)

keras_nlp/layers/transformer_decoder_test.py

@@ -332,7 +332,7 @@ def test_saved_model(self, save_format, filename):
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         model([decoder_sequence, encoder_sequence])
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        model.save(path, save_traces=False, save_format=save_format)
 
         loaded_model = keras.models.load_model(path)
         model_output = model([decoder_sequence, encoder_sequence])
@@ -358,7 +358,7 @@ def test_saved_model_without_cross_attention(self, save_format, filename):
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         model(decoder_sequence)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        model.save(path, save_traces=False, save_format=save_format)
         loaded_model = keras.models.load_model(path)
 
         model_output = model(decoder_sequence)

keras_nlp/layers/transformer_encoder_test.py

@@ -176,7 +176,7 @@ def test_saved_model(self, save_format, filename):
         data = tf.random.uniform(shape=[2, 4, 6])
         model_output = model(data)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        model.save(path, save_traces=False, save_format=save_format)
 
         loaded_model = keras.models.load_model(path)
         loaded_model_output = loaded_model(data)

keras_nlp/models/albert/albert_backbone_test.py

@@ -94,7 +94,9 @@ def test_error_for_invalid_num_groups(self):
     def test_saved_model(self, save_format, filename):
         model_output = self.backbone(self.input_batch)
         save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
+        self.backbone.save(
+            save_path, save_traces=False, save_format=save_format
+        )
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.

keras_nlp/models/albert/albert_classifier_test.py

@@ -126,7 +126,9 @@ def test_serialization(self):
     def test_saving_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
         save_path = os.path.join(self.get_temp_dir(), filename)
-        self.classifier.save(save_path, save_format=save_format)
+        self.classifier.save(
+            save_path, save_traces=False, save_format=save_format
+        )
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back

keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py

@@ -164,7 +164,7 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        model.save(path, save_traces=False, save_format=save_format)
 
         restored_model = keras.models.load_model(path)
         outputs = model(input_data)[0]["token_ids"]

keras_nlp/models/albert/albert_masked_lm_test.py

@@ -131,11 +131,13 @@ def test_classifier_fit_no_xla(self):
     def test_saved_model(self, save_format, filename):
         model_output = self.masked_lm.predict(self.raw_batch)
         save_path = os.path.join(self.get_temp_dir(), filename)
-        self.masked_lm.save(save_path, save_format=save_format)
+        self.masked_lm.save(
+            save_path, save_traces=False, save_format=save_format
+        )
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, AlbertMaskedLM)
         # Check that output matches.
         restored_output = restored_model.predict(self.raw_batch)
-        self.assertAllClose(model_output, restored_output)
+        self.assertAllClose(model_output, restored_output, atol=0.01, rtol=0.01)