Difficulty Interfacing TFP with Keras in TFP version 0.12.1 and TF version 2.4.1

With a recent update to tensorflow-probability and tensorflow I had to come up with workaround for a bug that presented in a model that I had been working on. In TF 2.3.0 and TFP 0.11.1 the following code worked:

def mixture_network(input_x,num_components):
    
    Noise = tf.keras.layers.GaussianNoise(ystd)
    
    mu =  Dense(num_components,name='mixture_means',activation='linear')(input_x) 
    sigma = Dense(num_components,name='mixture_std',activation=nnelu)(input_x) 
    p = Dense(num_components,name='mixing_weights',activation='softmax')(input_x)

    params = [mu,sigma,p]
    mixture = tfp.layers.DistributionLambda(make_distribution_fn=lambda params:tfd.MixtureSameFamily(
                mixture_distribution=tfd.Categorical(probs=params[2]),
                components_distribution=tfd.Normal(
                    loc=params[0],#tf.linspace(0.0,1.0,num_components),       
                    scale=params[1]#tf.ones(num_components)/(2*num_components)
                )
            )
    )(params)
    return Noise(mixture.prob(input_x)), params, mixture

With the following loss function:

def mixture_loss(y_true,y_pred):
    y_pred, params, _ = y_pred
    penalty = tf.pad(params[0][:,1:]-params[0][:,:-1],tf.constant([[0, 0,], [0, 1]]),
                     mode='CONSTANT',constant_values=2./num_components)
    penalty2 = params[1]
    penalty3 = tf.cast(tf.math.argmax(y_true)-tf.math.argmax(y_pred),dtype=tf.float32)
    return (1000.0*tf.math.reduce_mean(tf.abs(tf.math.log(y_true)-tf.math.log(y_pred)),axis=-1)
                               +1000*tf.math.reduce_mean(tf.abs(penalty-2./num_components),axis=-1)
                               +1000*tf.math.reduce_mean(tf.abs(penalty2-2./(3*num_components)),axis=-1)
                               +1000*tf.abs(penalty3)
           )

However, in TF version 2.4.1 and TFP version 0.12.1 I get the following error when trying to define the network:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-30-36de634f51f6> in <module>
     52     input_y = Input(shape=(length,)) #input Y values / true values from observation
     53 
---> 54     outputs = mixture_network_new(input_x,num_components)
     55     model = Model(inputs=[input_x,input_y],outputs=outputs)
     56     model.add_loss(mixture_loss6(input_y,outputs))

<ipython-input-30-36de634f51f6> in mixture_network_new(input_x, num_components)
     27     #return Noise(mixture.prob(input_x)), params, mixture
     28     #y = mixture.prob(input_x)
---> 29     return params, mixture, mixture.prob(input_x)
     30 def mixture_loss6(y_true,y_pred):
     31     params, _, input_x = y_pred

AttributeError: 'UserRegisteredTypeKerasTensor' object has no attribute 'prob'

Reproduce the provided mixture_network example with TensorFlow 2.4.1 and TensorFlow Probability 0.12.1, then compare it with the working TensorFlow 2.3.0 and TFP 0.11.1 versions. Start at the DistributionLambda output and the failing mixture.prob(input_x) call; done means identifying the compatibility regression and documenting a confirmed resolution or precise limitation.