Serialization of MixtureSameFamily layer fails

Hi,

I tried to save a model containing a MixtureSameFamily layer using tf.saved_model.save.
Unfortunately this fails: ValueError: Cannot pickle Tensor -- its value is not known statically.

Tf: 2.3.1
Tfp: 0.11.1

Any help would be appreciated.

Example

Code taken from MixtureSameFamily

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfpl = tfp.layers
tfk = tf.keras
tfkl = tf.keras.layers

# Load data -- graph of a [cardioid](https://en.wikipedia.org/wiki/Cardioid).
n = 2000
t = tfd.Uniform(low=-np.pi, high=np.pi).sample([n, 1])
r = 2 * (1 - tf.cos(t))
x = r * tf.sin(t) + tfd.Normal(loc=0., scale=0.1).sample([n, 1])
y = r * tf.cos(t) + tfd.Normal(loc=0., scale=0.1).sample([n, 1])

# Model the distribution of y given x with a Mixture Density Network.
event_shape = [1]
num_components = 5
params_size = tfpl.MixtureSameFamily.params_size(
    num_components,
    component_params_size=tfpl.IndependentNormal.params_size(event_shape))
model = tfk.Sequential([
  tfkl.Dense(12, activation='relu'),
  tfkl.Dense(params_size, activation=None),
  tfpl.MixtureSameFamily(num_components, tfpl.IndependentNormal(event_shape)),
])

# Fit.
batch_size = 100
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
              loss=lambda y, model: -model.log_prob(y))
          
model.fit(x, y,
          batch_size=batch_size,
          epochs=1,
          steps_per_epoch=n // batch_size)

tf.saved_model.save(model,"test_save")

Result

20/20 [==============================] - 0s 1ms/step - loss: 1.8857
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-9-817016578f49> in <module>
     36           steps_per_epoch=n // batch_size)
     37 
---> 38 tf.saved_model.save(model,"test_save")

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\saved_model\save.py in save(obj, export_dir, signatures, options)
    973   meta_graph_def = saved_model.meta_graphs.add()
    974 
--> 975   _, exported_graph, object_saver, asset_info = _build_meta_graph(
    976       obj, export_dir, signatures, options, meta_graph_def)
    977   saved_model.saved_model_schema_version = constants.SAVED_MODEL_SCHEMA_VERSION

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\saved_model\save.py in _build_meta_graph(obj, export_dir, signatures, options, meta_graph_def)
   1073         function_aliases[fdef.name] = alias
   1074 
-> 1075   object_graph_proto = _serialize_object_graph(saveable_view,
   1076                                                asset_info.asset_index)
   1077   meta_graph_def.object_graph_def.CopyFrom(object_graph_proto)

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\saved_model\save.py in _serialize_object_graph(saveable_view, asset_file_def_index)
    718 
    719   for obj, obj_proto in zip(saveable_view.nodes, proto.nodes):
--> 720     _write_object_proto(obj, obj_proto, asset_file_def_index,
    721                         saveable_view.function_name_map)
    722   return proto

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\saved_model\save.py in _write_object_proto(obj, proto, asset_file_def_index, function_name_map)
    759           version=versions_pb2.VersionDef(
    760               producer=1, min_consumer=1, bad_consumers=[]),
--> 761           metadata=obj._tracking_metadata)
    762       # pylint:enable=protected-access
    763     proto.user_object.CopyFrom(registered_type_proto)

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\keras\engine\base_layer.py in _tracking_metadata(self)
   3009   @property
   3010   def _tracking_metadata(self):
-> 3011     return self._trackable_saved_model_saver.tracking_metadata
   3012 
   3013   def _list_extra_dependencies_for_serialization(self, serialization_cache):

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\keras\saving\saved_model\base_serialization.py in tracking_metadata(self)
     52     # TODO(kathywu): check that serialized JSON can be loaded (e.g., if an
     53     # object is in the python property)
---> 54     return json_utils.Encoder().encode(self.python_properties)
     55 
     56   def list_extra_dependencies_for_serialization(self, serialization_cache):

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\keras\saving\saved_model\layer_serialization.py in python_properties(self)
     39   def python_properties(self):
     40     # TODO(kathywu): Add python property validator
---> 41     return self._python_properties_internal()
     42 
     43   def _python_properties_internal(self):

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\keras\saving\saved_model\model_serialization.py in _python_properties_internal(self)
     33 
     34   def _python_properties_internal(self):
---> 35     metadata = super(ModelSavedModelSaver, self)._python_properties_internal()
     36     # Network stateful property is dependent on the child layers.
     37     metadata.pop('stateful')

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\keras\saving\saved_model\layer_serialization.py in _python_properties_internal(self)
     57     )
     58 
---> 59     metadata.update(get_config(self.obj))
     60     if self.obj.input_spec is not None:
     61       # Layer's input_spec has already been type-checked in the property setter.

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\keras\saving\saved_model\layer_serialization.py in get_config(obj)
    116     # When loading, the program will attempt to revive the object from config,
    117     # and if that fails, the object will be revived from the SavedModel.
--> 118     config = generic_utils.serialize_keras_object(obj)['config']
    119 
    120   if config is not None:

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\keras\utils\generic_utils.py in serialize_keras_object(instance)
    243     name = get_registered_name(instance.__class__)
    244     try:
--> 245       config = instance.get_config()
    246     except NotImplementedError as e:
    247       if _SKIP_FAILED_SERIALIZATION:

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\keras\engine\sequential.py in get_config(self)
    463       # of `self.layers`). Note that `self._layers` is managed by the
    464       # tracking infrastructure and should not be used.
--> 465       layer_configs.append(generic_utils.serialize_keras_object(layer))
    466     config = {
    467         'name': self.name,

~\Anaconda3\envs\tf\lib\site-packages\tensorflow\python\keras\utils\generic_utils.py in serialize_keras_object(instance)
    243     name = get_registered_name(instance.__class__)
    244     try:
--> 245       config = instance.get_config()
    246     except NotImplementedError as e:
    247       if _SKIP_FAILED_SERIALIZATION:

~\Anaconda3\envs\tf\lib\site-packages\tensorflow_probability\python\layers\distribution_layer.py in get_config(self)
    294     """
    295     config = {
--> 296         'make_distribution_fn': _serialize_function(self._make_distribution_fn),
    297         'convert_to_tensor_fn': _serialize(self._convert_to_tensor_fn),
    298     }

~\Anaconda3\envs\tf\lib\site-packages\tensorflow_probability\python\layers\distribution_layer.py in _serialize_function(func)
   2045   pickler.dispatch_table[tf.Tensor] = _reduce_tensor
   2046 
-> 2047   pickler.dump(func)
   2048   return codecs.encode(buffer.getvalue(), 'base64').decode('ascii')
   2049 

~\Anaconda3\envs\tf\lib\site-packages\cloudpickle\cloudpickle_fast.py in dump(self, obj)
    561     def dump(self, obj):
    562         try:
--> 563             return Pickler.dump(self, obj)
    564         except RuntimeError as e:
    565             if "recursion" in e.args[0]:

~\Anaconda3\envs\tf\lib\site-packages\tensorflow_probability\python\layers\distribution_layer.py in _reduce_tensor(tensor)
   2024   val = tf.get_static_value(tensor)
   2025   if val is None:
-> 2026     raise ValueError('Cannot pickle Tensor -- '
   2027                      'its value is not known statically: {}.'.format(tensor))
   2028   return (tf.convert_to_tensor, (val,))

ValueError: Cannot pickle Tensor -- its value is not known statically: Tensor("MixtureSameFamily/independent_normal_6/IndependentNormal/Reshape:0", shape=(None, 5, 1), dtype=float32).

Start with tensorflow_probability/python/layers/distribution_layer.py, especially DistributionLambda.get_config and _serialize_function, then reproduce the failure with the supplied MixtureSameFamily example and tf.saved_model.save call. Done means the example saves the model without the reported non-static Tensor pickling error.