It is not possible to save a PixelCNN model

Code:

imgsize=64

dist = tfp.distributions.PixelCNN(
    image_shape=(imgsize, imgsize, 3),
    num_resnet=1,
    num_hierarchies=2,
    num_filters=160,
    num_logistic_mix=5,
    dropout_p=.3,
    high=1, low=0
)

image_input = layers.Input(shape=(imgsize, imgsize, 3))
log_prob = dist.log_prob(image_input)

model = tf.keras.Model(inputs=image_input, outputs=log_prob)
model.add_loss(-tf.reduce_mean(log_prob))

model.compile(
    optimizer=tf.keras.optimizers.Adam(.0001),
    metrics=[])

model_json = model.to_json()

Error:

NotImplementedError                       Traceback (most recent call last)
<ipython-input-12-d9399ed0ac2a> in <module>
----> 1 model_json = model.to_json()

~/anaconda3/envs/tf2-gpu/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/network.py in to_json(self, **kwargs)
   1252         A JSON string.
   1253     """
-> 1254     model_config = self._updated_config()
   1255     return json.dumps(
   1256         model_config, default=serialization.get_json_type, **kwargs)

~/anaconda3/envs/tf2-gpu/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/network.py in _updated_config(self)
   1230     from tensorflow.python.keras import __version__ as keras_version  # pylint: disable=g-import-not-at-top
   1231 
-> 1232     config = self.get_config()
   1233     model_config = {
   1234         'class_name': self.__class__.__name__,

~/anaconda3/envs/tf2-gpu/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/network.py in get_config(self)
    916     if not self._is_graph_network:
    917       raise NotImplementedError
--> 918     return copy.deepcopy(get_network_config(self))
    919 
    920   @classmethod

~/anaconda3/envs/tf2-gpu/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/network.py in get_network_config(network, serialize_layer_fn)
   1991           filtered_inbound_nodes.append(node_data)
   1992 
-> 1993     layer_config = serialize_layer_fn(layer)
   1994     layer_config['name'] = layer.name
   1995     layer_config['inbound_nodes'] = filtered_inbound_nodes

~/anaconda3/envs/tf2-gpu/lib/python3.6/site-packages/tensorflow_core/python/keras/utils/generic_utils.py in serialize_keras_object(instance)
    196 
    197   if hasattr(instance, 'get_config'):
--> 198     config = instance.get_config()
    199     serialization_config = {}
    200     for key, item in config.items():

~/anaconda3/envs/tf2-gpu/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py in get_config(self)
    497     # or that `get_config` has been overridden:
    498     if len(extra_args) > 1 and hasattr(self.get_config, '_is_default'):
--> 499       raise NotImplementedError('Layers with arguments in `__init__` must '
    500                                 'override `get_config`.')
    501     return config

NotImplementedError: Layers with arguments in `__init__` must override `get_config`.

Reproduce the example with tfp.distributions.PixelCNN, the Keras model construction, and model.to_json(). Trace the serialization path from model.to_json() into the PixelCNN-related layer configuration and verify that the model can be serialized without the reported NotImplementedError.