tuple has no attribute "rank" error when trying to build bayesian neural net

I'm trying to build a BNN but am encountering the error in the title. I tried to ensure I'm not passing a tuple to .shape.rank by :

• using the functional API with explicit Input (should make first denseflipout layer see a tensor, not tuple)
• I was told elsewhere that tfp 0.24.0 expects Keras 2, whereas my version of tf (2.19.0) installs Keras 3. So I upgraded to tfp-nightly 0.25.0 which is compatible with Keras 3 and still encounter the same issue.

Here's the mre with the full error I get, I'm using tensorflow 2.19.0 and tensorflow-probability 0.25.0. As you can see in the debug, "inp" is a tensor, not a tuple and should still be one when shape.rank is called.

import tensorflow as tf
import tensorflow_probability as tfp   # 0.25.0‑dev
tfpl = tfp.layers
inp = tf.keras.Input(shape=(10,))
print("[DEBUG] inp:", inp)
print("[DEBUG] type(inp):", type(inp))
print("[DEBUG] inp.shape:", inp.shape)
try:
    print("[DEBUG] inp.shape.rank:", inp.shape.rank)
except Exception as e:
    print("[DEBUG] could not read rank →", e)
print("[DEBUG] --------------------------------------------------")
out = tfpl.DenseFlipout(16)(inp)       # ← AttributeError: tuple has no attribute “rank”
tf.keras.Model(inp, out)

Output:

[DEBUG] inp: <KerasTensor shape=(None, 10), dtype=float32, sparse=False, ragged=False, name=keras_tensor_1>
[DEBUG] type(inp): <class 'keras.src.backend.common.keras_tensor.KerasTensor'>
[DEBUG] inp.shape: (None, 10)
[DEBUG] could not read rank → 'tuple' object has no attribute 'rank'
[DEBUG] --------------------------------------------------
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[2], line 13
     11     print("[DEBUG] could not read rank →", e)
     12 print("[DEBUG] --------------------------------------------------")
---> 13 out = tfpl.DenseFlipout(16)(inp)       # ← AttributeError: tuple has no attribute “rank”
     14 tf.keras.Model(inp, out)

File ~\AppData\Roaming\Python\Python39\site-packages\tf_keras\src\utils\traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File ~\AppData\Roaming\Python\Python39\site-packages\tf_keras\src\engine\input_spec.py:251, in assert_input_compatibility(input_spec, inputs, layer_name)
    244         raise ValueError(
    245             f'Input {input_index} of layer "{layer_name}" '
    246             "is incompatible with the layer: "
    247             f"expected max_ndim={spec.max_ndim}, "
    248             f"found ndim={ndim}"
    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(
    254             f'Input {input_index} of layer "{layer_name}" '
    255             "is incompatible with the layer: "
   (...)
    258             f"Full shape received: {tuple(shape)}"
    259         )

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

You create KerasTensor which don't have .shape.rank but only .ndim

You could have .shape.rank if you would create real Tensorflow's Tensor like this

t = tf.constant([[1, 2],[3, 4],[5, 6]], dtype=tf.float16)

print( t.shape.rank )
# 2
print( t ) 
# <tf.Tensor: shape=(3, 2), dtype=float16, numpy=array([[1., 2.],[3., 4.],[5., 6.]], dtype=float16)>