Cannot save the model when using Blockwise in a DistributionLambda layer

I have been trying to merge many output distributions into only one and then storing the model with no success. Currently, I got to the point that the model seems to work, although cannot store it. I added a DistributionLambda with a Blockwise distribution in order to merge them and then use log_prob in the loss, as shown in the code snippet below. When storing the model, it fails.

If there is any better choice to merge distributions into only one for using log_prob, it is greatly appreciated. Otherwise, the error seems to be related to #325 or other save-related bugs.

Here it is a code snippet, also uploaded as a colab here:

import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras import models
from tensorflow_probability import distributions
from tensorflow_probability import layers as tfp_layers


def likelihood_loss(y_true, y_pred):
    """Compute log likelihood losss."""
    # Shapes must match
    return -y_pred.log_prob(tf.reshape(y_true, tf.shape(y_pred)))


def distribution_fn(params):
    """Distribution function."""
    return distributions.Normal(params, 1.0)


def merge_distributions_fn(params):
    """Merge distribution functions."""
    return distributions.Blockwise(params)


x_data = np.array([[0.0], [1.0], [2.0], [3.0], [4.0]])
y_data = np.array([[1.0, 2.0], [2.0, 3.0], [3.0, 4.0], [4.0, 5.0], [5.0, 6.0]])

print(x_data.shape)
print(y_data.shape)

input_layer = layers.Input(shape=x_data.shape[1:])
distr1_layer = tfp_layers.DistributionLambda(
    make_distribution_fn=distribution_fn)(input_layer)
distr2_layer = tfp_layers.DistributionLambda(
    make_distribution_fn=distribution_fn)(input_layer)
last_layer = tfp_layers.DistributionLambda(
    make_distribution_fn=merge_distributions_fn)([distr1_layer, distr2_layer])

model = models.Model(inputs=[input_layer], outputs=[last_layer])
model.compile(loss=likelihood_loss, optimizer='Adam')
print(model.summary())
model.fit(x=x_data, y=y_data, epochs=2)
model.save('test')

The exception error message:

TypeError: Tensor("inputs:0", shape=(None, 1), dtype=float32) must be either `tfd.Distribution`-like or `callable`.

Start with the provided Colab and code snippet, reproducing the failure at model.save('test') after building the DistributionLambda and Blockwise layers. Trace how these distributions are handled during model serialization; done means the shown model can be saved without the reported TypeError.