`TypeSpec` issue using `JointDistribution` with `DistributionLambda` layer

[psilentp]

I'm getting an error trying to build a Keras model (using either the sequential or functional API) using a DistributionLambda output layer with a lambda function that returns JointDistribution as the make_distribution_fn.

The error is: ValueError: KerasTensor only supports TypeSpecs that have a shape field; got TensorTupleSpec, which does not have a shape..

TFP version 0.16.0 TF version 2.8.0

Here is test case (which was originally in tensorflow_probability/python/layers/distribution_layer_test.py but removed with a810afe).

x = tf.keras.Input(shape=())
y = tfp.layers.VariableLayer(shape=[2, 4, 3], dtype=tf.float32)(x)
y = tf.keras.layers.Dense(5, use_bias=False)(y)
y = tfp.layers.DistributionLambda(
    lambda t: tfd.JointDistributionSequential([  # pylint: disable=g-long-lambda
        tfd.Gamma(t[..., 0], t[..., 1]),
        tfd.Normal(t[..., 2], 1),
        lambda m, s: tfd.Normal(loc=m, scale=s),
    ]),
)(y)
m = tf.keras.Model(x, y)

I haven't been able to work out why the test was removed, but was wondering if there might be a new approach to accommodate the same end.