How to use the edward2.layers function in edward2

recurrent_regularizer=make_regularizer(),
            bias_initializer=bias_initializer,
            bias_regularizer=make_regularizer(bias_uncertainty))
      else:
        lstm_cell = tf.keras.layers.LSTMCell(
            rnn_dim,
            kernel_regularizer=tf.keras.regularizers.l2(l2),
            recurrent_regularizer=tf.keras.regularizers.l2(l2),
            bias_regularizer=tf.keras.regularizers.l2(l2))
      cells.append(lstm_cell)
    self.rnn_layer = tf.keras.layers.RNN(cells, return_sequences=False)

    # 2. Affine layer on combination of RNN output and context features.
    if self.hidden_layer_dim > 0:
      if hidden_uncertainty:
        self.hidden_layer = ed.layers.DenseReparameterization(
            self.hidden_layer_dim,
            activation=tf.nn.relu6,
            kernel_initializer="trainable_he_normal",
            kernel_regularizer=make_regularizer(),
            bias_initializer=bias_initializer,
            bias_regularizer=make_regularizer(bias_uncertainty))
      else:
        self.hidden_layer = tf.keras.layers.Dense(
            self.hidden_layer_dim,
            activation=tf.nn.relu6,
            kernel_regularizer=tf.keras.regularizers.l2(l2),
            bias_regularizer=tf.keras.regularizers.l2(l2))

    # 3. Output layer.
    self.output_uncertainty = output_uncertainty
    if self.output_uncertainty: