how to train neural network mixture model with tfp

Hello, I am trying to write a simple gaussian mixture model using tfp. Basically I want to implement

an MDN with tensorflow probability, but i have a mistake

`from keras.models import Sequential

from keras.layers.core import Dense, Dropout
from keras.callbacks import History
from keras.layers.recurrent import LSTM
from keras.models import model_from_json

c = 1 #The number of outputs we want to predict
m = 2 #The number of distributions we want to use in the mixture

from keras import backend as K

i=tf.keras.layers.Input(shape=((14,40)))
x=tf.keras.layers.LSTM(64, return_sequences=True)(i)
#x=Dropout(0.5)(x)
x=tf.keras.layers.LSTM(64, return_sequences=True)(x)
#x=Dropout(0.5)(x)

#x=Dropout(0.5)(x)
x=tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(64))(x)
x=tf.keras.layers.AveragePooling1D()(x)
x=tf.keras.layers.Flatten()(x)
mu = tf.keras.layers.Dense(int(c*m), name='mdn_mus')(x)#(self.inputs)
sigma = tf.keras.layers.Dense(m, activation=K.softplus, name='mdn_sigmas')(x)#(self.inputs)
xi = tf.keras.layers.Dense(m, activation=K.softplus, name='mdn_xi')(x)#(self.inputs)
pi = tf.keras.layers.Dense(m, activation=K.softmax, name='mdn_pi')(x)#(self.inputs)

mdn_out = tf.keras.layers.concatenate([mu, sigma,xi, pi], name='mdn_out')
#mdn_out=Dense(1, activation=K.softmax)(mdn_out)

model = tf.keras.Model(inputs=i, outputs=mdn_out )
print(model.summary())
optim = tf.keras.optimizers.Adam(lr=0.01)
model.compile(optimizer=optim, loss=mdn_loss)

model.fit(X_unrolled_train, y_train, epochs=5, batch_size=64,
shuffle=False, validation_data=(X_unrolled_test, y_test))`

and the loss function

def mdn_loss(y_true, y_pred):
out_mu, out_sigma,out_xi, out_pi = tf.split(y_pred, num_or_size_splits=[m, m, m,m],axis=-1, name='mdn_coef_split')
mus = tf.split(out_mu, num_or_size_splits=m, axis=1)
sigs = tf.split(out_sigma, num_or_size_splits=m, axis=1)
xis = tf.split(out_xi, num_or_size_splits=m, axis=1)
cat = tfd.Categorical(logits=out_pi)
coll = [tfd.Normal(loc=loc, scale=scale) for loc, scale in zip(mus, sigs)]
mixture = tfd.Mixture(cat=cat, components=coll)
loss = mixture.log_prob(y_true)
loss = tf.negative(loss)
loss = tf.reduce_mean(loss)
return loss

i have this error

ValueError: components[0] batch shape must be compatible with cat shape >and other component batch shapes

Start in the shown mdn_loss function and inspect the shapes produced by tf.split, the Categorical logits, and each Normal component. Verify that the category and component batch shapes are compatible, then run the model's fit call until mixture.log_prob(y_true) completes without the ValueError.