TF2, Tensorflow Probability random seed generator and VAE

Playing around with Variational Autoencoders for some days. I am trying to fit a small toy function with a small model.

I first implemented the model using the Keras Functional API, with the following code:

def define_tfp_encoder(latent_dim, n_inputs=2, kl_weight=1):
    prior = tfd.MultivariateNormalDiag(loc=tf.zeros(latent_dim))
    
    input_x = Input((n_inputs,))
    input_c = Input((1,))
    
    dense = Dense(25, activation='relu', name='tfpenc/dense_1')(input_x)
    dense = Dense(32, activation='relu', name='tfpenc/dense_2')(dense)
    
    dense_z_params = Dense(tfpl.MultivariateNormalTriL.params_size(latent_dim), name='tfpenc/z_params')(dense)
    dense_z = tfpl.MultivariateNormalTriL(latent_dim, name='tfpenc/z')(dense_z_params)
                    #activity_regularizer=tfpl.KLDivergenceRegularizer(prior) # weight=kl_weight
                        
    
    kld = tfpl.KLDivergenceAddLoss(prior, name='tfpenc/kld_add')(dense_z)
    model = Model(inputs=input_x, outputs=kld)
    
    return model

def define_tfp_decoder(latent_dim, n_inputs=2):
    input_c = Input((1,), name='tfpdec/cond_input')
    input_n = Input((latent_dim,))
    
    dense = Dense(15, activation='relu', name='tfpdec/dense_1')(input_n)
    dense = Dense(32, activation='relu', name='tfpdec/dense_2')(dense)
    dense = Dense(tfpl.IndependentNormal.params_size(n_inputs), name='tfpdec/output')(dense)
    output = tfpl.IndependentNormal((n_inputs,))(dense)

    model = Model(input_n, output)

    return model

def get_custom_unconditional_vae():
    latent_size = 5
    encoder = define_tfp_encoder(latent_dim=latent_size)
    decoder = define_tfp_decoder(latent_dim=latent_size)
    
    encoder.trainable = True
    decoder.trainable = True
    x = encoder.input
    z = encoder.output
    out = decoder(z)
    
    vae = Model(inputs=x, outputs=out)
    vae.compile(loss=lambda x, pred: -pred.log_prob(x), optimizer='adam')    
    return encoder, decoder, vae

The vae-model was then fitted and trained on 3000 epochs. However, it only produced garbage for a very simple quadratic function to fit.

Now it comes: When creating the exact same model using the sequential API it works as expected and the desired function gets approximated nicely:

And it becomes even stranger for me: After running tf.random.set_seed(None) the model created using the Functional API also works as expected - What am I missing or not understanding correctly so far? - I assume that there are some differences regarding tf.random.set_seed when using the Sequential vs. the Functional API but... ?

Thanks in advance, codax