def probabilistic_pca(D, K, N, sigma):
    """
    :param D: number of data dimensions
    :param K: number of latent factors
    :param N: number of data points
    :param sigma: noise parameter
    """
    w = sample("w", Normal(zeros(D, K), ones(D, K)))
    z = sample("z", Normal(zeros(K, N), ones(K, N)))
    x = sample("x", Normal(w @ z, sigma * ones(D, N)))

    return x, w, z

def probabilistic_pca(D, K, N, sigma):
    """
    :param D: number of data dimensions
    :param K: number of latent factors
    :param N: number of data points
    :param sigma: noise parameter
    """
    w = ed.Normal(loc=tf.zeros([D, K]),
                  scale=tf.ones([D, K]),
                  name="w")
    z = ed.Normal(loc=tf.zeros([K, N]),
                  scale=tf.ones([K, N]),
                  name="z")
    x = ed.Normal(loc=tf.matmul(w, z),
                  scale=sigma * tf.ones([D, N]),
                  name="x")

    return x, w, z