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 (acompaniying notebook here - https://github.com/taboola/mdn-tensorflow-notebook-example/blob/master/mdn-tf2.ipynb).

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]