ValueError: Cell is empty with saving a model with tfpl.IndependentBernoulli layer

Hi,

I'm encountering ValueError: Cell is empty when trying to save a model with tfpl.IndependentBernoulli layer, with tensorflow version 2.4.1 and tensorflow_probability version 0.12.1. Here is a simple example
`
import sys
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers as L
import tensorflow_probability as tfp
tfpl = tfp.layers
print(sys.version, tf.version, tfp.version)

3.6.9 (default, Jan 26 2021, 15:33:00)

[GCC 8.4.0] 2.4.1 0.12.1

test_model = keras.Sequential([
L.Dense(2, input_shape=(10,)),
tfpl.IndependentBernoulli(2)
])
test_model.save('test_model')
which outputs: ---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
in
3 tfpl.IndependentBernoulli(2)
4 ])
----> 5 test_model.save('test_model')

~/.local/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py in save(self, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)
2000 # pylint: enable=line-too-long
2001 save.save_model(self, filepath, overwrite, include_optimizer, save_format,
-> 2002 signatures, options, save_traces)
2003
2004 def save_weights(self,

~/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/save.py in save_model(model, filepath, overwrite, include_optimizer, save_format, signatures, options, save_traces)
155 else:
156 saved_model_save.save(model, filepath, overwrite, include_optimizer,
--> 157 signatures, options, save_traces)
158
159

~/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/saved_model/save.py in save(model, filepath, overwrite, include_optimizer, signatures, options, save_traces)
87 with distribution_strategy_context._get_default_replica_context(): # pylint: disable=protected-access
88 with utils.keras_option_scope(save_traces):
---> 89 save_lib.save(model, filepath, signatures, options)
90
91 if not include_optimizer:

~/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/save.py in save(obj, export_dir, signatures, options)
1031
1032 _, exported_graph, object_saver, asset_info = _build_meta_graph(
-> 1033 obj, signatures, options, meta_graph_def)
1034 saved_model.saved_model_schema_version = constants.SAVED_MODEL_SCHEMA_VERSION
1035

~/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/save.py in _build_meta_graph(obj, signatures, options, meta_graph_def)
1196
1197 with save_context.save_context(options):
-> 1198 return _build_meta_graph_impl(obj, signatures, options, meta_graph_def)

~/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/save.py in _build_meta_graph_impl(obj, signatures, options, meta_graph_def)
1161
1162 object_graph_proto = _serialize_object_graph(saveable_view,
-> 1163 asset_info.asset_index)
1164 meta_graph_def.object_graph_def.CopyFrom(object_graph_proto)
1165

~/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/save.py in _serialize_object_graph(saveable_view, asset_file_def_index)
753 for obj, obj_proto in zip(saveable_view.nodes, proto.nodes):
754 _write_object_proto(obj, obj_proto, asset_file_def_index,
--> 755 saveable_view.function_name_map)
756 return proto
757

~/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/save.py in _write_object_proto(obj, proto, asset_file_def_index, function_name_map)
798 version=versions_pb2.VersionDef(
799 producer=1, min_consumer=1, bad_consumers=[]),
--> 800 metadata=obj._tracking_metadata)
801 # pylint:enable=protected-access
802 proto.user_object.CopyFrom(registered_type_proto)

~/.local/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py in _tracking_metadata(self)
3077 @property
3078 def _tracking_metadata(self):
-> 3079 return self._trackable_saved_model_saver.tracking_metadata
3080
3081 def _list_extra_dependencies_for_serialization(self, serialization_cache):

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

~/.local/lib/python3.6/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):

~/.local/lib/python3.6/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')

~/.local/lib/python3.6/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.

~/.local/lib/python3.6/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:

~/.local/lib/python3.6/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:

~/.local/lib/python3.6/site-packages/tensorflow/python/keras/engine/sequential.py in get_config(self)
469 # of self.layers). Note that self._layers is managed by the
470 # tracking infrastructure and should not be used.
--> 471 layer_configs.append(generic_utils.serialize_keras_object(layer))
472 config = {
473 'name': self.name,

~/.local/lib/python3.6/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:

~/.local/lib/python3.6/site-packages/tensorflow_probability/python/layers/distribution_layer.py in get_config(self)
785 'validate_args': self._validate_args
786 }
--> 787 base_config = super(IndependentBernoulli, self).get_config()
788 return dict(list(base_config.items()) + list(config.items()))
789

~/.local/lib/python3.6/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 }

~/.local/lib/python3.6/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

~/.local/lib/python3.6/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]:

/usr/lib/python3.6/pickle.py in dump(self, obj)
407 if self.proto >= 4:
408 self.framer.start_framing()
--> 409 self.save(obj)
410 self.write(STOP)
411 self.framer.end_framing()

/usr/lib/python3.6/pickle.py in save(self, obj, save_persistent_id)
474 f = self.dispatch.get(t)
475 if f is not None:
--> 476 f(self, obj) # Call unbound method with explicit self
477 return
478

~/.local/lib/python3.6/site-packages/cloudpickle/cloudpickle_fast.py in save_function(self, obj, name)
743 else:
744 return self._save_reduce_pickle5(
--> 745 *self._dynamic_function_reduce(obj), obj=obj
746 )
747

~/.local/lib/python3.6/site-packages/cloudpickle/cloudpickle_fast.py in _save_reduce_pickle5(self, func, args, state, listitems, dictitems, state_setter, obj)
680 self.save_reduce(
681 func, args, state=None, listitems=listitems,
--> 682 dictitems=dictitems, obj=obj
683 )
684 # backport of the Python 3.8 state_setter pickle operations

/usr/lib/python3.6/pickle.py in save_reduce(self, func, args, state, listitems, dictitems, obj)
608 else:
609 save(func)
--> 610 save(args)
611 write(REDUCE)
612

/usr/lib/python3.6/pickle.py in save(self, obj, save_persistent_id)
474 f = self.dispatch.get(t)
475 if f is not None:
--> 476 f(self, obj) # Call unbound method with explicit self
477 return
478

/usr/lib/python3.6/pickle.py in save_tuple(self, obj)
749 write(MARK)
750 for element in obj:
--> 751 save(element)
752
753 if id(obj) in memo:

/usr/lib/python3.6/pickle.py in save(self, obj, save_persistent_id)
474 f = self.dispatch.get(t)
475 if f is not None:
--> 476 f(self, obj) # Call unbound method with explicit self
477 return
478

/usr/lib/python3.6/pickle.py in save_tuple(self, obj)
734 if n <= 3 and self.proto >= 2:
735 for element in obj:
--> 736 save(element)
737 # Subtle. Same as in the big comment below.
738 if id(obj) in memo:

/usr/lib/python3.6/pickle.py in save(self, obj, save_persistent_id)
474 f = self.dispatch.get(t)
475 if f is not None:
--> 476 f(self, obj) # Call unbound method with explicit self
477 return
478

~/.local/lib/python3.6/site-packages/dill/_dill.py in save_cell(pickler, obj)
1175 def save_cell(pickler, obj):
1176 log.info("Ce: %s" % obj)
-> 1177 f = obj.cell_contents
1178 pickler.save_reduce(_create_cell, (f,), obj=obj)
1179 log.info("# Ce")

ValueError: Cell is empty`

I'm a bit surprised that this bug is not seen, since Bernoulli distribution is very commonly used. Please help fix it.

Start with the minimal Sequential example and the test_model.save('test_model') entry point, then follow the traceback through IndependentBernoulli.get_config and _serialize_function. Done means saving this model no longer raises ValueError: Cell is empty; add or update a regression test for the reproducible case if the repository has a relevant test location.