Bug when saving model with MixtureLogistic Layer

[henrypinkard]

When trying to save a keras model with a a tfp.layers.MixtureLogistic layer, the following error occurs:

WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.

---------------------------------------------------------------------------
OperatorNotAllowedInGraphError            Traceback (most recent call last)
Input In [5], in <cell line: 12>()
      8 out = tfp.layers.MixtureLogistic(9, [1,])(x)
     10 model = tfk.Model(inputs, out)
---> 12 tfk.models.save_model(model, '/path/to/save/dir/' + 'test')

File ~/miniconda3/envs/info_theory/lib/python3.9/site-packages/keras/utils/traceback_utils.py:67, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     65 except Exception as e:  # pylint: disable=broad-except
     66   filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67   raise e.with_traceback(filtered_tb) from None
     68 finally:
     69   del filtered_tb

File ~/miniconda3/envs/info_theory/lib/python3.9/contextlib.py:126, in _GeneratorContextManager.__exit__(self, typ, value, traceback)
    124 if typ is None:
    125     try:
--> 126         next(self.gen)
    127     except StopIteration:
    128         return False

File ~/miniconda3/envs/info_theory/lib/python3.9/site-packages/tensorflow_probability/python/layers/distribution_layer.py:220, in DistributionLambda.__call__(self, inputs, *args, **kwargs)
    218 def __call__(self, inputs, *args, **kwargs):
    219   self._enter_dunder_call = True
--> 220   distribution, _ = super(DistributionLambda, self).__call__(
    221       inputs, *args, **kwargs)
    222   self._enter_dunder_call = False
    223   return distribution

OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

This can be reproduced with the example below:

import tensorflow.keras as tfk
import tensorflow_probability as tfp

input_shape = (28, )

inputs = tfk.Input(shape=input_shape)
x = tfk.layers.Dense(9)(inputs)
out = tfp.layers.MixtureLogistic(9, [1,])(x)

model = tfk.Model(inputs, out)

tfk.models.save_model(model, '/path/to/save/dir/' + 'test')

It's definitely coming from the MixtureLogistic layer, because saving the output of the Dense directly works fine

tensorflow version 2.8.0 tensorflow_probability version: 0.16.0

[csuter]

Can you try saving with the argument save_format='h5'? This works for me.

You will also need to change your Dense layer size I think. Use params_size to automatically get the right number of params:

x = tfk.layers.Dense(tfp.layers.MixtureLogistic.params_size(9, [1]))(inputs)

[henrypinkard]

Yes this seems to work for me, thanks!