Bug in automated transformations for variables with simplex support

Hi everyone,

I’m trying to use the automated transformations functionality that was recently added to Edward to perform inference using KLqp for a model with a Dirichlet prior. Unfortunately, for variables with support under the simplex, such as a Dirichlet, the automated transformation appears to be broken and results in an error from TensorFlow/Edward.

See below for a very simple script which illustrates the issue. I get the following error from this script: “ValueError: Dimensions must be equal, but are 3 and 4 for ‘inference/sample/Dirichlet/log_prob/mul’ (op: ‘Mul’) with input shapes: [3], [4].” I’m using the latest development version of Edward, along with TensorFlow 1.3.0. Any ideas on how to resolve this issue?

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import edward as ed
import tensorflow as tf

from edward.models import Dirichlet

alpha = [1.0, 2.0, 3.0]
phi = Dirichlet(alpha)

inference = ed.KLqp([phi])
qphi = inference.latent_vars[phi]

If I had to guess, it’s because Dirichlet has simplex support, which means that the automatically transformed variational model will have an additional dimension (hence the “3 and 4”).

However, the way ADVI (what the auto transform implements) should work is to create that 4-dim prior, transform the provided prior to the unconstrained space, then match the transformed prior to an unconstrained variational model (which wouldn’t have this error). So potentially there is an issue with the implementation.

I’m going to open an issue on the repo.

Here is the github issue for reference.