AttributeError: 'tuple' object has no attribute 'rank' with tfp.layers.DenseVariational and Keras Input (TF 2.18.0, TFP 0.24.0)

Describe the bug
I am encountering an AttributeError: 'tuple' object has no attribute 'rank' when using tfp.layers.DenseVariational as the first layer following a tf.keras.Input layer in a Keras Functional API model. The error seems to occur during Keras's internal InputSpec compatibility check, where it expects a tf.TensorShape object but receives a Python tuple for the input shape from the DenseVariational layer.

To Reproduce
Steps to reproduce the behavior:

1. Environment: Python 3.11.9 on Windows 10.
2. Installed packages:
   • tensorflow==2.18.0
   • tensorflow-probability==0.24.0
   • numpy==1.26.4
3. Run the following minimal code snippet:

import tensorflow as tf
import tensorflow_probability as tfp

tfk = tf.keras
tfkl = tf.keras.layers
tfd = tfp.distributions

print(f"TensorFlow Version: {tf.__version__}")
print(f"TensorFlow Probability Version: {tfp.__version__}")
print(f"NumPy Version: {np.__version__}")

input_feature_size = 60 
num_classes = 3
num_train_samples_example = 1000

kl_divergence_weight = 1.0 / num_train_samples_example

try:
    inputs = tfk.Input(shape=(input_feature_size,), name="input_layer")
    
    posterior_fn_constructor = tfp.layers.default_mean_field_normal_fn(is_singular=True)
    prior_fn_constructor = tfp.layers.default_multivariate_normal_fn

    x = tfp.layers.DenseVariational(
        units=64, 
        make_posterior_fn=posterior_fn_constructor,
        make_prior_fn=prior_fn_constructor,
        kl_weight=kl_divergence_weight,
        activation='relu',
        name="dv1"
    )(inputs) 
    
    outputs = tfkl.Dense(units=num_classes, name="logits")(x)
    
    model = tfk.Model(inputs=inputs, outputs=outputs, name="bnn_test_model")
    model.summary() 
    print("Model created successfully.")

except Exception as e:
    print(f"\nError during model creation: {e}")
    import traceback
    traceback.print_exc()



An unexpected error occurred in main: 'tuple' object has no attribute 'rank'
Traceback (most recent call last):
  File "C:\Users\W10\PycharmProjects\eurusd5\eurusd5.py", line 239, in <module>
    accuracies, true_labels, predicted_labels = walk_forward_tfp_bnn(
                                                ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\W10\PycharmProjects\eurusd5\eurusd5.py", line 151, in walk_forward_tfp_bnn
    bnn_model = create_bnn_model(input_shape_for_model, num_classes, len(X_train_scaled))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\W10\PycharmProjects\eurusd5\eurusd5.py", line 78, in create_bnn_model
    x = tfp.layers.DenseVariational(
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\W10\PycharmProjects\eurusd5\.venv\Lib\site-packages\tf_keras\src\utils\traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\W10\PycharmProjects\eurusd5\.venv\Lib\site-packages\tf_keras\src\engine\input_spec.py", line 251, in assert_input_compatibility
    ndim = x.shape.rank
           ^^^^^^^^^^^^
AttributeError: 'tuple' object has no attribute 'rank'



System information
OS Platform and Distribution: Windows 10
TensorFlow installed from: pip
TensorFlow version: 2.18.0
TensorFlow Probability version: 0.24.0
Python version: 3.11.9
NumPy version: 1.26.4
CUDA/cuDNN version: N/A (CPU only)
GPU model and memory: N/A (CPU only)
Additional context
The issue appears to be specific to how tfp.layers.DenseVariational (in TFP 0.24.0) interacts with the Keras Functional API's input shape handling when using TensorFlow 2.18.0, particularly for 1D vector inputs. Attempts to use batch_input_shape in tf.keras.Input or explicitly pass input_shape to DenseVariational did not resolve this specific "rank" error. The error occurs even when using TFP's default make_posterior_fn and make_prior_fn.
Thank you for your help!

Reproduce the minimal Functional API example with tfp.layers.DenseVariational, then inspect the compatibility check in tf_keras/src/engine/input_spec.py and the DenseVariational call path named in the traceback. Compare the input shape handling that reaches x.shape.rank; done means the model can be created and summarized without the tuple-rank AttributeError.