Remove keras.experimental.export_saved_model API. It was deprecated in TF 1.15, which was 3 years ago. It was never a stable API in the first place (it was an experimental API) and had no backwards compatibility guarantee.

Consider upgrading to:

1. `model.save()`
2. `tf.saved_model.save()`

PiperOrigin-RevId: 474040485

Author: fchollet

tensorflow_probability/python/layers/distribution_layer_test.py

@@ -351,28 +351,6 @@ def assertSerializable(self, model, batch_size=1):
     self.assertAllEqual(self.evaluate(model(x).log_prob(y)),
                         self.evaluate(model_copy(x).log_prob(y)))
 
-  def assertExportable(self, model, batch_size=1):
-    """Assert a Keras model supports export_saved_model/load_from_saved_model.
-
-    Args:
-      model: A Keras model with Tensor output.
-      batch_size: The batch size to use when checking that the model produces
-        the same results as a serialized/deserialized copy.  Default value: 1.
-    """
-    batch_shape = [batch_size]
-
-    input_shape = batch_shape + model.input.shape[1:].as_list()
-    dtype = model.input.dtype.as_numpy_dtype
-
-    model_dir = self.create_tempdir()
-    tf1.keras.experimental.export_saved_model(model, model_dir.full_path)
-    model_copy = tf1.keras.experimental.load_from_saved_model(
-        model_dir.full_path)
-
-    x = np.random.uniform(-3., 3., input_shape).astype(dtype)
-    self.assertAllEqual(self.evaluate(model(x)), self.evaluate(model_copy(x)))
-    self.assertAllEqual(model.predict(x), model_copy.predict(x))
-
   def test_serialization(self):
     model = tfk.Sequential([
         tfkl.Dense(2, input_shape=(5,)),
@@ -382,15 +360,6 @@ def test_serialization(self):
     ])
     self.assertSerializable(model)
 
-    model = tfk.Sequential([
-        tfkl.Dense(2, input_shape=(5,)),
-        # pylint: disable=g-long-lambda
-        tfpl.DistributionLambda(lambda t: tfd.Normal(
-            loc=t[..., 0:1], scale=tf.exp(t[..., 1:2]))),
-        tfkl.Lambda(lambda d: d.mean() + d.stddev())
-    ])
-    self.assertExportable(model, batch_size=4)
-
   @staticmethod
   def _make_distribution(t):
     return tfpl.MixtureSameFamily.new(t, 3, tfpl.IndependentNormal([2]))
@@ -405,16 +374,6 @@ def test_serialization_static_method(self):
     model.compile(optimizer='adam', loss='mse')
     self.assertSerializable(model, batch_size=3)
 
-    model = tfk.Sequential([
-        tfkl.Dense(15, input_shape=(5,)),
-        tfpl.DistributionLambda(
-            DistributionLambdaSerializationTest._make_distribution,
-            convert_to_tensor_fn=tfd.Distribution.mean),
-        tfkl.Lambda(lambda x: x + 1.)
-    ])
-    model.compile(optimizer='adam', loss='mse')
-    self.assertExportable(model)
-
   def test_serialization_closure_over_lambdas_tensors_and_numpy_array(self):
     if six.PY2 and not tf.executing_eagerly():
       self.skipTest('Serialization of constant graph-mode Tensors is not '
@@ -433,14 +392,6 @@ def make_distribution(t):
     ])
     self.assertSerializable(model, batch_size=4)
 
-    model = tfk.Sequential([
-        tfkl.Dense(15, input_shape=(5,)),
-        # pylint: disable=unnecessary-lambda
-        tfpl.DistributionLambda(lambda t: make_distribution(t)),
-        tfkl.Lambda(lambda d: d.mean() + d.stddev())
-    ])
-    self.assertExportable(model, batch_size=2)
-
 
 @test_util.test_graph_and_eager_modes
 class DistributionLambdaVariableCreation(test_util.TestCase):
@@ -874,34 +825,6 @@ def test_serialization(self):
     self.assertAllEqual(self.evaluate(model(x).log_prob(ones)),
                         self.evaluate(model_copy(x).log_prob(ones)))
 
-  def test_model_export(self):
-    event_shape = [3, 2]
-    params_size = self.layer_class.params_size(event_shape)
-    batch_shape = [4]
-
-    low = self._build_tensor(-3., dtype=self.dtype)
-    high = self._build_tensor(3., dtype=self.dtype)
-    x = self.evaluate(tfd.Uniform(low, high).sample(
-        batch_shape + [params_size], seed=42))
-
-    model = tfk.Sequential([
-        tfkl.Dense(params_size, input_shape=(params_size,), dtype=self.dtype),
-        self.layer_class(event_shape, validate_args=True,
-                         convert_to_tensor_fn='mean', dtype=self.dtype),
-        # NOTE: For TensorFlow to be able to serialize the graph (i.e., for
-        # serving), the model must output a Tensor and not a Distribution.
-        tfkl.Dense(1, dtype=self.dtype),
-    ])
-    model.compile(optimizer='adam', loss='mse')
-
-    model_dir = self.create_tempdir()
-    tf1.keras.experimental.export_saved_model(model, model_dir.full_path)
-    model_copy = tf1.keras.experimental.load_from_saved_model(
-        model_dir.full_path)
-
-    self.assertAllEqual(self.evaluate(model(x)), self.evaluate(model_copy(x)))
-    self.assertEqual(self.dtype, model(x).dtype.as_numpy_dtype)
-
 
 @test_util.test_graph_and_eager_modes
 class _IndependentBernoulliTest(_IndependentLayerTest):
@@ -1246,35 +1169,6 @@ def test_serialization(self):
     self.assertAllEqual(self.evaluate(model(x).log_prob(ones)),
                         self.evaluate(model_copy(x).log_prob(ones)))
 
-  def test_model_export(self):
-    n = 5
-    event_shape = [3, 2]
-    params_size = self.layer_class.params_size(n, event_shape)
-    batch_shape = [4]
-
-    low = self._build_tensor(-3., dtype=self.dtype)
-    high = self._build_tensor(3., dtype=self.dtype)
-    x = self.evaluate(tfd.Uniform(low, high).sample(
-        batch_shape + [params_size], seed=42))
-
-    model = tfk.Sequential([
-        tfkl.Dense(params_size, input_shape=(params_size,), dtype=self.dtype),
-        self.layer_class(n, event_shape, validate_args=True,
-                         convert_to_tensor_fn='mean', dtype=self.dtype),
-        # NOTE: For TensorFlow to be able to serialize the graph (i.e., for
-        # serving), the model must output a Tensor and not a Distribution.
-        tfkl.Dense(1, dtype=self.dtype),
-    ])
-    model.compile(optimizer='adam', loss='mse')
-
-    model_dir = self.create_tempdir()
-    tf1.keras.experimental.export_saved_model(model, model_dir.full_path)
-    model_copy = tf1.keras.experimental.load_from_saved_model(
-        model_dir.full_path)
-
-    self.assertAllEqual(self.evaluate(model(x)), self.evaluate(model_copy(x)))
-    self.assertEqual(self.dtype, model(x).dtype.as_numpy_dtype)
-
 
 @test_util.test_graph_and_eager_modes
 class _MixtureLogisticTest(_MixtureLayerTest):