The shape of a probabilistic layer throws and error

[shifdz]

I am a building model with TensorFlow probability layers. When I do, model.output.shape, I get an error:

AttributeError: 'UserRegisteredSpec' object has no attribute '_shape'

If I do, output_shape = tf.shape(model.output) it gives a Keras Tensor:

<KerasTensor: shape=(5,) dtype=int32 inferred_value=[None, 3, 128, 128, 128] (created by layer 'tf.compat.v1.shape_15') 

How can I get the actual values [None, 3, 128, 128, 128]? I tried output_shape.get_shape(), but that gives the Tensor shape [5].

code to reproduce:

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd

tfd = tfp.distributions

model = tf.keras.Sequential()
model.add(tf.keras.layers.Input(10))

model.add(tf.keras.layers.Dense(2, activation="linear"))
model.add(
    tfp.layers.DistributionLambda(
        lambda t: tfd.Normal(
            loc=t[..., :1], scale=1e-3 + tf.math.softplus(0.1 * t[..., 1:])
        )
    )
)
model.compile(
    optimizer=tf.keras.optimizers.Adam(),
    loss="mean_absolute_error",
    # List of metrics to monitor
    metrics="mean_absolute_error",
)
model.save("tf_test_model.h5")

model = load_model("tf_test_model.h5")
model.output.shape

[Frightera]

I think the workaround is the use _inferred_value on that KerasTensor:

tf.shape(model.output)._inferred_value should give the desired value.