title: How to clone/duplicate a TensorFlow Probability neural network model
body: <p>I have a TensorFlow Probability model that is built similar to models described in this <a href="https://www.youtube.com/watch?v=BrwKURU-wpk" rel="nofollow noreferrer">YouTube Video</a>.  Here's the code to build the model:</p>
  <pre><code>def posterior_mean_field(kernel_size: int, bias_size: int, dtype: Any) -&gt; tf.keras.Model:
      n = kernel_size + bias_size
      c = np.log(np.expm1(1.))
  
      return tf.keras.Sequential([
          tfp.layers.VariableLayer(2 * n, dtype=dtype),
          tfp.layers.DistributionLambda(lambda t: tfd.Independent(tfd.Normal(loc=t[..., :n],
                                                                             scale=1e-5 + tf.nn.softplus(c + t[..., n:])),
                                                                  reinterpreted_batch_ndims=1)),
      ])
  
  
  def prior_trainable(kernel_size: int, bias_size: int, dtype: Any) -&gt; tf.keras.Model:
      n = kernel_size + bias_size
      return tf.keras.Sequential([
          tfp.layers.VariableLayer(n, dtype=dtype),
          tfp.layers.DistributionLambda(lambda t: tfd.Independent(
              tfd.Normal(loc=t, scale=1),
              reinterpreted_batch_ndims=1)),
      ])
  
  
  def build_model():
      model = keras.Sequential([
          tfp.layers.DenseVariational(64, activation='relu', input_shape=[len(train_dataset.keys())],
                                      make_posterior_fn=posterior_mean_field,
                                      make_prior_fn=prior_trainable),
          layers.Dense(64, activation='relu'),
          layers.Dense(1),
      ])
      optimizer = tf.keras.optimizers.RMSprop(0.001)
      model.compile(loss='mse', optimizer=optimizer, metrics=['mae', 'mse'])
      return model
  
  
  model = build_model()
  model.build((3, 10))
  </code></pre>
  <p>When I remove the TensorFlow Probability layer (1st layer) in the model, I can clone the model and copy its weights like this:</p>
  <pre><code>import copy
  from tensorflow.keras.models import clone_model
  model_weights = copy.deepcopy(model.get_weights())
  model_copy = clone_model(model)
  model_copy.set_weights(model_weights)
  </code></pre>
  <p>However, when the TensorFlow Probability layer is present I get this error:</p>
  <pre><code>Traceback (most recent call last):
    File &quot;/Users/jisom/opt/miniconda3/envs/ic-hours/lib/python3.8/site-packages/IPython/core/interactiveshell.py&quot;, line 3398, in run_code
      exec(code_obj, self.user_global_ns, self.user_ns)
    File &quot;&lt;ipython-input-6-349eb0e7c1e5&gt;&quot;, line 1, in &lt;cell line: 1&gt;
      model_new = clone_model(model)
    File &quot;/Users/jisom/opt/miniconda3/envs/ic-hours/lib/python3.8/site-packages/keras/models.py&quot;, line 448, in clone_model
      return _clone_sequential_model(
    File &quot;/Users/jisom/opt/miniconda3/envs/ic-hours/lib/python3.8/site-packages/keras/models.py&quot;, line 326, in _clone_sequential_model
      if isinstance(layer, InputLayer) else layer_fn(layer))
    File &quot;/Users/jisom/opt/miniconda3/envs/ic-hours/lib/python3.8/site-packages/keras/models.py&quot;, line 56, in _clone_layer
      return layer.__class__.from_config(layer.get_config())
    File &quot;/Users/jisom/opt/miniconda3/envs/ic-hours/lib/python3.8/site-packages/keras/engine/base_layer.py&quot;, line 727, in get_config
      raise NotImplementedError('Layer %s has arguments in `__init__` and '
  NotImplementedError: Layer DenseVariational has arguments in `__init__` and therefore must override `get_config`.
  </code></pre>
  <p>I can see some information about how to deal with this error in <a href="https://stackoverflow.com/questions/58678836/notimplementederror-layers-with-arguments-in-init-must-override-get-conf">this StackOverflow question</a>, but it in that question there's a custom-built transformer class that can be modified.  I'm trying to use the <code>clone_model</code> function in keras, which I don't directly control.  How can I clone/duplicate a model, including its weights, if the model includes TensorFlow Probability layers?</p>