Tensorflow Probability Glow has duplicated variable names

Hi,

I am playing with the default Glow model implemented in tensorflow_probability and I find that many variables have duplicated names in the graph (e.g. many variables are called 'Variable:0'). That prevents, for example, saving the model in an h5 file.

Am I doing something wrong? Reproducible code below

import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_probability as tfp
import math
tfd = tfp.distributions
tfb = tfp.bijectors

class GlowModel(tf.keras.Model):
    def __init__(self, input_shape, **kwargs):
        super(GlowModel, self).__init__(**kwargs)
        self.glow = tfb.Glow(output_shape=input_shape,
                        coupling_bijector_fn=tfb.GlowDefaultNetwork,
                        exit_bijector_fn=tfb.GlowDefaultExitNetwork,
                        num_glow_blocks=3,
                        num_steps_per_block=32)
        z_shape = self.glow.inverse_event_shape(input_shape)

        self.pz = tfd.Sample(tfd.Normal(0., 1.), z_shape)

        self.px = self.glow(self.pz)

    def call(self, inputs, training=None, **kwargs):
        return self.px
    
model = GlowModel(input_shape=(32, 32, 1))
print(model.trainable_variables)

Thanks

Start by running the reproducible GlowModel example and inspecting model.trainable_variables. Read tfb.Glow, tfb.GlowDefaultNetwork, and tfb.GlowDefaultExitNetwork to determine where variable names are assigned; done means the model has distinct variable names and can be saved to an H5 file.