How to get the shape of a probabilistic layer?

Question

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 error:

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.output.shape

Answer 1

tf.shape will return a KerasTensor which is not easy to get the output shape directly.

However you can do this:

tf.shape(model.output)
>> `<KerasTensor: shape=(2,) dtype=int32 inferred_value=[None, 1] (created by layer 'tf.compat.v1.shape_168')>`

You want to get inferred_value, so:

tf.shape(model.output)._inferred_value
>> [None, 1]

Basically you can access any layer's output shape with:

tf.shape(model.layers[idx].output)._inferred_value

where idx is the index of the desired layer.