Probability Output Mixture distributions

Question

I am trying to create a regression model with a probability output layer in tensorflow. I would like to be able to use a specified mixture of distributions in the output layer.

For example, from the tensorflow documentation the example that is given for tfp.layers.MixtureSameFamily is:

num_components = 3
model = tfk.Sequential([
  tfkl.Dense(12, activation='relu'),
  tfkl.Dense(params_size, activation=None),
  tfpl.MixtureSameFamily(num_components, tfpl.IndependentNormal(event_shape)),
])

I would like to do a similar thing but instead of using as mixture of, for example 3 normal distributions, use 1 truncated normal distribution followed by 2 normal distributions.

I have tried to make this work by using tfp.distributions.Mixture in a similar way to tfp.layers.MixtureSameFamily but not sure how the implementation would work as they inherit from different classes.

perhaps the best way is to write a new class that inherits from tfp.layers.MixtureSameFamily but has methods that can use a mixture of different distributions?

Any ideas on how a probability output layer combining different distributions could work would be appreciated.

Thanks.