Cannot save model with stacked distribution layers as a saved model (h5 works)

I've another model saving issue that to me seems slightly different from others reported, this is on python 3.6.9 with:

tensorflow==2.4.1
tensorflow-probability==0.12.1

Saving models with stacked distribution layers throws an exception, unless I save as an h5 (that model can be loaded and behaves as expected). This seems to be related to the handling of the _TensorCoercible behaving differently when tracing the graph during a save, leading to the output of a distribution layer being coerced to a tensor. Here's a minimal example to reproduce:

import tensorflow as tf
from tensorflow_probability.python.bijectors import Reshape
from tensorflow_probability.python.distributions import MultivariateNormalDiag
from tensorflow_probability.python.layers import DistributionLambda

inputs = tf.keras.Input(shape=(1, 20))
x = DistributionLambda(
    lambda mu: MultivariateNormalDiag(loc=mu)
)(inputs)

x = DistributionLambda(
    lambda d: Reshape(event_shape_in=[20], event_shape_out=[2, 10])(d)
)(x)

model = tf.keras.Model(inputs, x)
print(model.summary())

model.save("model.h5")  # Works fine
model.save("model")  # Throws a TypeError

The traceback is a monster, but here's the tail:

  File ".../python3.6/site-packages/tensorflow/python/eager/function.py", line 3206, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File ".../python3.6/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File ".../python3.6/site-packages/tensorflow/python/eager/def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File ".../python3.6/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 527, in wrapper
    ret = method(*args, **kwargs)
  File ".../python3.6/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 171, in wrap_with_training_arg
    lambda: replace_training_and_call(False))
  File ".../python3.6/site-packages/tensorflow/python/keras/utils/control_flow_util.py", line 115, in smart_cond
    pred, true_fn=true_fn, false_fn=false_fn, name=name)
  File ".../python3.6/site-packages/tensorflow/python/framework/smart_cond.py", line 54, in smart_cond
    return true_fn()
  File ".../python3.6/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 170, in <lambda>
    training, lambda: replace_training_and_call(True),
  File ".../python3.6/site-packages/tensorflow/python/keras/saving/saved_model/utils.py", line 167, in replace_training_and_call
    return wrapped_call(*args, **kwargs)
  File ".../python3.6/site-packages/tensorflow/python/keras/saving/saved_model/save_impl.py", line 570, in call_and_return_conditional_losses
    call_output = layer_call(inputs, *args, **kwargs)
  File ".../python3.6/site-packages/tensorflow_probability/python/layers/distribution_layer.py", line 252, in call
    inputs, *args, **kwargs)
  File ".../python3.6/site-packages/tensorflow/python/keras/layers/core.py", line 917, in call
    result = self.function(inputs, **kwargs)
  File ".../python3.6/site-packages/tensorflow_probability/python/layers/distribution_layer.py", line 179, in _fn
    convert_to_tensor_fn=maybe_composite_convert_to_tensor_fn)
  File ".../python3.6/site-packages/tensorflow_probability/python/layers/internal/distribution_tensor_coercible.py", line 54, in __new__
    distribution, type(distribution)))
TypeError: `distribution` argument must be a `tfd.Distribution` instance; saw "Tensor("reshape_1/forward/Reshape:0", shape=(None, None, 2, 10), dtype=float32)" of type "<class 'tensorflow.python.framework.ops.Tensor'>".

Start by running the stacked DistributionLambda reproduction and compare model.h5 with the SavedModel path. Trace the save call through tensorflow/python/keras/saving/saved_model/save_impl.py and tensorflow_probability/python/layers/distribution_layer.py into distribution_tensor_coercible.py. Done means the stacked distribution model saves and reloads successfully without the TensorCoercible type error.