Models using a KLDivergenceRegularizer cannot be saved

I've been working through the TFP VAE tutorial, and tried to save the model after training. However, this was not possible in SavedModel format.

Command:

model.save(saved_model_path, save_format="tf", save_traces=True)

Produces the following error:

` File "/Users/xxx/venvs/ml_venv/lib/python3.9/site-packages/tensorflow_probability/python/layers/distribution_layer.py", line 1251, in call *
return self._kl_divergence_fn(distribution_a)

File "/Users/xxx/venvs/ml_venv/lib/python3.9/site-packages/tensorflow_probability/python/layers/distribution_layer.py", line 1370, in _fn  **  
    kl = kl_divergence_fn(distribution_a, distribution_b_)   

File "/Users/xxx/venvs/ml_venv/lib/python3.9/site-packages/tensorflow_probability/python/layers/distribution_layer.py", line 1354, in kl_divergence_fn. 
    distribution_a.log_prob(z) - distribution_b.log_prob(z),  
AttributeError: 'Tensor' object has no attribute 'log_prob'`. 

Command:

model.save(saved_model_path, save_format="tf", save_traces=False)

Produces the following error:

<tensorflow_probability.python.layers.distribution_layer.KLDivergenceRegularizer object at 0x16b79c5b0> does not implement get_config()

Indeed it does not, KLDivergenceRegularizer does not have a get_config implementation.

Sidenote: I cannot load from h5 because MultivariateNormalTriL requires a positional argument:
TypeError: __init__() missing 1 required positional argument: 'event_size'

Versions:
Python 3.9.13
tensorflow-probability==0.17.0
tensorflow==2.9.0

This is also discussed in #742 but I'm starting a new issue to make it more prominent.

Start in tensorflow_probability/python/layers/distribution_layer.py, especially KLDivergenceRegularizer and the serialization paths exercised by model.save. Reproduce the issue with save_traces=True and save_traces=False using the VAE tutorial setup. Done means the model can be saved and loaded in SavedModel format without either reported error.