How to train a mixture of gaussians bayesian neural network?

Question

I am trying to train a mixture neural network, but I'm getting the following error:

AssertionError: Exception encountered when calling layer "model_16" (type Functional).

Could not compute output KerasTensor(type_spec=TensorSpec(shape=(None,), dtype=tf.float32, name=None), name='distribution_lambda/tensor_coercible/value/Mixture/sample/Reshape:0', description="created by layer 'distribution_lambda'")

I suspect the error is related to the shapes of my tensors. The model compiles, and I can inspect the structure of my network. It is as follows:

The code to reproduce the previous image and the error is the following:

from tensorflow.keras import initializers
import tensorflow_probability as tfp
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
from keras import backend as K
from keras.models import Model

tfpl = tfp.layers
tfkl = tf.keras.layers
tfd = tfp.distributions


inputs = keras.Input(shape=(8,), name="my_inputs")

x = layers.Dense(8, activation="relu", name="dense_1",
                    kernel_initializer=initializers.RandomNormal(stddev=0.02),
    bias_initializer=initializers.Zeros())(inputs)

x0 = layers.Dense(1, activation="linear", name="dense_loc1")(x)
x1 = layers.Dense(1, activation="linear", name="dense_loc2")(x)
x2 = layers.Dense(1, activation="relu", name="dense_scale1")(x)
x3 = layers.Dense(1, activation="relu", name="dense_scale2")(x)
x4 = layers.Dense(2, activation="softmax", name="dense_softmax")(x)

#x1 = tf.keras.layers.Reshape((-1, None))(x1)
x0 = tf.reshape(x0, [-1], name="reshape_loc1")
x1 = tf.reshape(x1, [-1], name="reshape_loc2")
x2 = tf.reshape(x2, [-1], name="reshape_scale1")
x3 = tf.reshape(x3, [-1], name="reshape_scale2")

x2 = 1e-3+tf.math.abs(x2, name="abs_scale1")
x3 = 1e-3+tf.math.abs(x3, name="abs_scale2")


outputs = tfpl.DistributionLambda(lambda  t : tfd.Mixture(
        cat=tfd.Categorical(probs= t[4]),
        components=[tfd.Normal(loc=t[0], scale=t[2]),
                    tfd.Normal(loc=t[1], scale=t[3])])
                                )([x0, x1, x2, x3, x4])

model_sc = keras.Model(inputs=inputs, outputs=outputs)

negloglik = lambda y_true, y_pred: -y_pred.log_prob(y_true)

model_sc.compile(optimizer=tf.optimizers.Adam(learning_rate=0.0000005), loss=negloglik)

model_sc.summary()
My model
#Early Stopping Callback
callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=600, min_delta=0, mode='auto', baseline=None, restore_best_weights=True)

#Testing with sample data

model = model_sc  # My model
x_data = np.array([[ 1.0052079e+00,  3.6000001e-01,  2.7777779e-01, -4.8919746e-01, -5.7142860e-01,  1.4186612e+00,  0.0000000e+00,  5.1869620e-12]])


for layer in model.layers:
    layer_name = layer.name
    print(layer_name)
    
    intermediate_layer_model = Model(inputs=model.input,
                                     outputs=model.get_layer(layer_name).output)
    intermediate_output = intermediate_layer_model.predict([x_data] )
    print(tf.shape(intermediate_output))
    print(intermediate_output)