Fitting an arbitrary number of component distributions using tfp.layers.MixtureSameFamily

Hey everyone, I was just wondering if there is a simple a way to train a network to produce a Gaussian Mixture Model (GMM) with an arbitrary number of component distributions using tfp.layers.MixtureSameFamily? I'm looking to produce a probabilistic model, specifically a GMM, that best fits some given input and output training data by using a neural network. Implementing the code to train for a fixed number of component distributions is simple enough and something I have already done, however, my overall goal is to eventually implement a neural network that can generate a GMM with an arbitrary number of component distributions (aka modes) that is "necessary" for each test case. Here is my code for a fixed number of component distributions:

x = some 2-d input data
y = some 2-d output data

event_shape = y.shape[-1]
num_components = 5
params_size = tfp.layers.MixtureSameFamily.params_size(
num_components,
component_params_size=tfp.layers.MultivariateNormalTriL.params_size(event_shape))

model = tf.keras.Sequential([
tf.keras.layers.Dense(params_size),
tfp.layers.MixtureSameFamily(num_components, tfp.layers.MultivariateNormalTriL(event_shape)),
])

model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01),
loss=lambda y, model1: -model1.log_prob(y),
metrics=[])
batch_size = 250
model.fit(x, y,
batch_size=batch_size,
epochs=500,
steps_per_epoch=n // batch_size,
verbose=False,
shuffle=True)

yhat_mix=model(x_tst)

Is there a simple way to modify this code so that the number of component distributions is arbitrary? Or would it require a major overhaul of the actual code implementation? I thought making "num_components" a trainable parameter might work, however, the labels in the training data don't actually contain any information regarding the number of modes so I don't think training the network in the traditional sense would work in this case. Alternatively, I've considered just using a very large number of component distributions (1000+) and hoping the network naturally produces only a select few modes with non-zero covariance. This seems extremely crude though and would probably be quite computationally expensive.

I'm not really sure on how else approach this problem as I'm quite new to tensorflow and python in general, so any help or ideas would be greatly appreciated!

Start by reading the tfp.layers.MixtureSameFamily API and the fixed-component example in the issue, focusing on how num_components and params_size are used. A complete result would define and implement a supported way to choose the number of components from training data, with documented behavior and validation for the resulting model.