I am trying to fit a Mixture Density Network model using tensorflow, following the example posted in Taboola’s blog - Predicting Probability Distributions Using Neural Networks 2 (acompaniying notebook here - https://github.com/taboola/mdn-tensorflow-notebook-example/blob/master/mdn-tf2.ipynb 1).
The data I am fitting are positive integers so I initially modified the blog example to use a Poisson distribution instead of a Normal from the blog example:
def poisson_mdn_cost(lambda_, y_):
    dist = tfp.distributions.Poisson(rate = lambda_)
    return tf.reduce_mean(-dist.log_prob(y_))

epochs = 500
batch_size = 50
learning_rate = 0.0003
InputLayer = Input(shape=(1,))
Layer_1 = Dense(12,activation="tanh")(InputLayer)
Layer_2 = Dense(12,activation="tanh")(Layer_1)
lambda_ = Dense(1, activation=lambda x: tf.nn.elu(x) + 1)(Layer_2)

Y = Input(shape=(1,))
lossF =  poisson_mdn_cost(lambda_,Y)
model2 = Model(inputs=[InputLayer, Y], outputs=[lambda_])
model2.add_loss(lossF)
model2.compile(optimizer = 'adam',
               metrics = ['mse'],
             )
history_cache2 = model2.fit([x_data, y_data], 
                            verbose=0, 
                            epochs=epochs,
                            batch_size=batch_size)
print('Final cost: {0:.4f}'.format(history_cache2.history['loss'][-1]))

The Poisson MDN fitted fine and I am now running an experiment with a Beta-binomial MDN using the following code:
def betabin_mdn_cost( a_,b_, y):
   dist = tfp.distributions.BetaBinomial(total_count =10,
                                         concentration0=a_,
                                         concentration1=b_) 
   return tf.reduce_mean(-dist.log_prob(y))

epochs = 500
batch_size = 50
learning_rate = 0.0003
InputLayer = Input(shape=(1,))
Layer_1 = Dense(12,activation="tanh")(InputLayer)
Layer_2 = Dense(12,activation="tanh")(Layer_1)
a = Dense(1, activation="relu")(Layer_2)
b = Dense(1, activation="relu")(Layer_2)

Y = Input(shape=(1,))
lossF =  betabin_mdn_cost(a,b,Y)
model3 = Model(inputs=[InputLayer, Y], outputs=[a,b])
model3.add_loss(lossF)
model3.compile(optimizer = 'adam',
               metrics = ['mse'],
             )
history_cache3 = model3.fit([x_data, y_data], 
                            verbose=0, 
                            epochs=epochs,
                            batch_size=batch_size)
print('Final cost: {0:.4f}'.format(history_cache3.history['loss'][-1]))

The data is truncated between 0 and 10, so fixed the Beta-binomial total_count parameter to 10. As far as I can see, the latter code has minimal changes compared with the Poisson above, essentially just swapping the Poisson cost function by the Beta-binomial and adding the additional parameter. However this time I can’t get it to run and get the error:
  TypeError: Cannot convert a symbolic Keras input/output to a numpy array. This error may indicate that you're trying to pass a symbolic value to a NumPy call, which is not supported. Or, you may be trying to pass Keras symbolic inputs/outputs to a TF API that does not register dispatching, preventing Keras from automatically converting the API call to a lambda layer in the Functional Model.

I an new to tf and have exhausted my ability to try to get this to work. I Was wondering if there is anybody out there with experience in fitting these types of models using tf, that might have some insights?
[tf.version.VERSION = 2.6.0]