Multiple features in Mixture Density Networks Tutorial

In the MDN example for Edward, if I have > 1 feature (lets say D=2 in the example, where it’s currently 1 ), I get this error when I run the following section of code in the example. I had modified the dataset to be of shape (N,2) instead of the (N,1) in the tutorial.

X_ph = tf.placeholder(tf.float32, [None, D])
y_ph = tf.placeholder(tf.float32, [None, D])

def neural_network(X):
  """loc, scale, logits = NN(x; theta)"""
  # 2 hidden layers with 15 hidden units
  net = tf.layers.dense(X, 15, activation=tf.nn.relu)
  net = tf.layers.dense(net, 15, activation=tf.nn.relu)
  locs = tf.layers.dense(net, K, activation=None)
  scales = tf.layers.dense(net, K, activation=tf.exp)
  logits = tf.layers.dense(net, K, activation=None)
  return locs, scales, logits

locs, scales, logits = neural_network(X_ph)

cat = Categorical(logits=logits)
components = [Normal(loc=loc, scale=scale) for loc, scale
              in zip(tf.unstack(tf.transpose(locs)),
y = Mixture(cat=cat, components=components, value=tf.zeros_like(y_ph))


The error is:

ValueError                                Traceback (most recent call last)
<ipython-input-9-31ddf9990b8f> in <module>()
     16               in zip(tf.unstack(tf.transpose(locs), axis=0),
     17                      tf.unstack(tf.transpose(scales), axis=0))]
---> 18 y = Mixture(cat=cat, components=components, value=tf.zeros_like(y_ph))
     19 # Note: A bug exists in Mixture which prevents samples from it to have
     20 # a shape of [None]. For now fix it using the value argument, as

~/anaconda3/lib/python3.6/site-packages/edward/models/ in __init__(self, *args, **kwargs)
     19     # to use _candidate's docstring, must write a new __init__ method
     20     def __init__(self, *args, **kwargs):
---> 21       _RandomVariable.__init__(self, *args, **kwargs)
     22     __init__.__doc__ = _candidate.__init__.__doc__
     23     _params = {'__doc__': _candidate.__doc__,

~/anaconda3/lib/python3.6/site-packages/edward/models/ in __init__(self, *args, **kwargs)
    121         raise ValueError(
    122             "Incompatible shape for initialization argument 'value'. "
--> 123             "Expected %s, got %s." % (expected_shape, value_shape))
    124       else:
    125         self._value = t_value

ValueError: Incompatible shape for initialization argument 'value'. Expected (?,), got (?, 2)


It appears that the D=2 is not transmitted to the network - I am new to Edward, and I would appreciate any suggestions on fixing this.