How to use the edward2.layers.DenseReparameterization function in edward2

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:
      self.output_layer = ed.layers.DenseReparameterization(
          output_layer_dim,
          kernel_regularizer=make_regularizer(),
          bias_initializer=bias_initializer,
          bias_regularizer=make_regularizer(bias_uncertainty))
    else:
      self.output_layer = tf.keras.layers.Dense(
          output_layer_dim,
          kernel_regularizer=tf.keras.regularizers.l2(l2),
          bias_regularizer=tf.keras.regularizers.l2(l2))