Как перевести этот код Эдварда в код TFP?

Я закодировал модель вероятностной матричной факторизации в Эдварде. Я пытаюсь перенести его на TFP, но я не уверен, как определить термины логарифмической вероятности и расхождения KL. Вот код в Эдварде -

# MODEL
U = Normal(
    loc=0.0,
    scale=1.0,
    sample_shape=[n_latent_dims, batch_size])
V = Normal(
    loc=0.0,
    scale=1.0,
    sample_shape=[n_latent_dims, n_features])
R = Bernoulli(logits=tf.matmul(tf.transpose(U), V))
R_ph = tf.placeholder(tf.int32, [None, n_features])

# INFERENCE
qU = Normal(
    loc=tf.get_variable("qU/loc",
                        [n_latent_dims, batch_size]),
    scale=tf.nn.softplus(
        tf.get_variable("qU/scale",
                        [n_latent_dims, batch_size])))
qV = Normal(
    loc=tf.get_variable("qV/loc",
                        [n_latent_dims, n_features]),
    scale=tf.nn.softplus(
        tf.get_variable("qV/scale",
                        [n_latent_dims, n_features])))
qR = Bernoulli(logits=tf.matmul(tf.transpose(qU), qV))
qR_avg = Bernoulli(
    logits=tf.reduce_mean(qR.parameters['logits'], 0))
log_likli = tf.reduce_mean(qR_avg.log_prob(R_ph), 1)
inference = ed.KLqp(
    {
        U: qU,
        V: qV
    }, data={self.R: self.R_ph})
inference_init = inference.initialize()
init = tf.global_variables_initializer()