DistributionLambda mixed Distribution and Tensor inputs

Hey, I'm wondering if it's somehow possible to pass mixed Distribution and Tensor inputs to the DistributionLambda layer. Here's an example of a fail case:

import tensorflow as tf
import tensorflow_probability as tfp

tfk = tf.keras
tfkl = tf.keras.layers
tfpl = tfp.layers
tfd = tfp.distributions
tfb = tfp.bijectors


def main():
    output_shape = (5, )
    conditions_input = tfkl.Input(shape=(3, ))

    latent_distribution = tfd.Independent(
        tfd.Normal(loc=tf.zeros(output_shape), scale=tf.ones(output_shape)),
        reinterpreted_batch_ndims=1)

    shift_and_log_scale_model = tfk.Sequential((
        tfkl.Dense(16, activation='relu'),
        tfkl.Dense(
            tfpl.IndependentNormal.params_size(output_shape), activation=None),
    ))

    shift_and_log_scale = shift_and_log_scale_model(conditions_input)

    def transform_distribution(inputs):
        shift_and_log_scale, latent_distribution = inputs
        shift, log_scale_diag = tf.split(shift_and_log_scale, 2, axis=-1)
        bijector = tfb.Chain((
            tfb.Affine(
                shift=shift,
                scale_diag=tf.exp(log_scale_diag)),
            tfb.Tanh(),
        )[::-1])

        transformed_distribution = tfd.TransformedDistribution(
            distribution=latent_distribution,
            bijector=bijector,
            batch_shape=shift.shape[:1])

        return transformed_distribution

    transformed_distribution_model = tfk.Sequential((
        tfpl.DistributionLambda(
            make_distribution_fn=transform_distribution),
    ))

    transformed_distribution = transformed_distribution_model((
        shift_and_log_scale, latent_distribution))

    assert isinstance(transformed_distribution, tfd.Distribution)


if __name__ == '__main__':
    main()

This code fails with

ValueError: Cannot convert a partially known TensorShape to a Tensor: (?,)

which I believe is caused by the fact that the shapes of shift_and_log_scale and latent_distribution don't match (which is understandable since they are of different type).

My main question is: should this work or is what I'm doing here weird and could be achieved some other way?

For the context, I'm trying to build a model which transforms a simple latent distribution (e.g. unit Gaussian) into a new distribution via some parameterized (learnable) bijector (e.g. affine transformation to produce a new Gaussian or RealNVP transformation to produce more complicated RealNVP distribution). I know that I could build the distribution in the above example case simply by:

transformed_distribution = tfk.Sequential((
    tfkl.Dense(16, activation='relu'),
    tfkl.Dense(
        tfpl.IndependentNormal.params_size(output_shape), activation=None),
    tfpl.IndependentNormal(output_shape),
))(shift_and_log_scale)

However, the reason I'm trying to separate the latent distribution and the transformation is that I'd like to be able to control the sampling (for example do exponential smoothing) in the latent space. Effectively, I would like to do something like the following pseudo-code does, except with keras models:

conditions_np =  tf.random.uniform((1, 3))
latent_samples = smooth(latent_distribution.sample(1)) # Exponential smoothing
transformed_latent_sample = tf.Session.run(transformed_distribution.bijector.forward(latent_samples)[0], feed_dict={conditions_input: conditions_np}) # Not sure how this would be done with the models

It's totally possible that I'm trying to fit the keras models and tfp distributions together in a way that's not meant to be done. I'm still figuring out what the best practice for building trainable distributions like the one above is, so all feedback is welcome.

Start by running the reported example and inspect the DistributionLambda entry point, especially how it handles the mixed shift_and_log_scale and latent_distribution inputs and their shapes. Determine whether mixed Distribution and Tensor inputs are intended to work; done means either supporting this use case with a passing reproduction or documenting the supported approach for parameterized transformations and latent-space sampling.