Something wrong with the creation of a Keras model with probabilistic layers.

I am applying the use of Keras models that utilize probabilistic layers. The model below, which used to work normally, now presents the following error:

inputs = tf.keras.Input(shape=(X_train.shape[1],))
x = layers.Dense(500, activation='relu')(inputs)
x2 = layers.Dense(350, activation='relu')(x)
x3 = layers.Dense(250, activation='relu')(x2)
x4 = layers.Dense(150, activation='relu')(x3)
x5 = layers.Dense(50, activation='relu')(x4)
x6 = layers.Dense(15, activation='relu')(x5)
y_variational = tfpl.DenseVariational(4, make_prior_fn=prior, make_posterior_fn=posterior, kl_weight=1/X_train.shape[0], activation='tanh')(x6)

distribution_params = layers.Dense(units=216*2)(y_variational)
outputs = tfpl.IndependentNormal(216)(distribution_params)

PBNN3 = tf.keras.Model(inputs=inputs, outputs=outputs)

AttributeError Traceback (most recent call last)
in <cell line: 8>()
6 x5 = layers.Dense(50, activation='relu')(x4) #15
7 x6 = layers.Dense(15, activation='relu')(x5)
----> 8 y_variational = tfpl.DenseVariational(4, make_prior_fn=prior, make_posterior_fn=posterior, kl_weight=1/X_train.shape[0], activation='tanh')(x6)
9 #y_variational = layers.Dense(10, activation='relu')(x)
10 distribution_params = layers.Dense(units=216*2)(y_variational)

1 frames
/usr/local/lib/python3.10/dist-packages/tf_keras/src/engine/input_spec.py in assert_input_compatibility(input_spec, inputs, layer_name)
249 )
250 if spec.min_ndim is not None:
--> 251 ndim = x.shape.rank
252 if ndim is not None and ndim < spec.min_ndim:
253 raise ValueError(

AttributeError: 'tuple' object has no attribute 'rank'

Something is not working with TensorFlow version 2.17.0 and TensorFlow Probability version 0.24. Still works for Tensorflow: 2.15.1 and TensorFlow Probability: 0.22.1.

Thank you!

Start by reproducing the model construction shown in the issue with TensorFlow 2.17.0 and TensorFlow Probability 0.24, then compare it with TensorFlow 2.15.1 and TensorFlow Probability 0.22.1, which are reported to work. Trace the input compatibility failure at tf_keras/src/engine/input_spec.py:251; done means identifying the version interaction and confirming a compatible behavior or a targeted fix.