tf.keras.Model.predict() shape ValueError with explicit convert_to_tensor_fn lambda in tfp.layers.MultivariateNormalTriL

Consider the following TF 2.1 and TFP 0.9 untrained model that outputs multiple (sample_shape=3) samples from MultivariateNormalTriL:

import tensorflow as tf
import tensorflow_probability as tfp

d = 2
model = tf.keras.Sequential([
    tf.keras.layers.Dense(units=tfp.layers.MultivariateNormalTriL.params_size(d), 
                          input_shape=(d,)),
    tfp.layers.MultivariateNormalTriL(event_size=d, 
                                      convert_to_tensor_fn=lambda s: s.sample(sample_shape=3)),
])

print(model.predict(tf.zeros(shape=(1, d))))

The model definition follows examples in distribution_layer_test.py. The sample code should work, but it raises:

ValueError: Mismatch between expected batch size and model output batch size. Output shape = (3, 1, 2), expected output shape = shape (1, 1, 2)

Why does tf.keras.Model.predict() expect sample_shape==1? Is this a bug, or is the model misspecified somehow? How should a TFP layer be defined to pass multiple samples downstream? Thank you.

Start with the provided TF 2.1/TFP 0.9 reproduction and the referenced distribution_layer_test.py, focusing on MultivariateNormalTriL and the explicit convert_to_tensor_fn. Trace how tf.keras.Model.predict validates the output shape against the input batch, then determine the supported behavior for multiple samples and capture it in a regression test or clarified documentation.