MDN Implementation using Edward2

[Niknafs]

The current example on MDN from Edward tutorials needs small modifications to run on edward2. Documentation covering these modifications will be appreciated.

[dustinvtran]

This would be hugely useful if anyone's interested in contributing!

For the record, the following model works:

def build_neural_network():
  inputs = tf.keras.layers.Input(...)
  net = tf.keras.layers.Dense(X, 15, activation='relu')(inputs)
  net = tf.keras.layers.Dense(15, activation=tf.nn.relu)(net)
  locs = tf.keras.layers.Dense(K, activation=None)(net)
  scales = tf.keras.layers.Dense(K, activation=tf.exp)(net)
  logits = tf.keras.layers.Dense(K, activation=None)(net)
  model = tf.keras.Model(inputs=inputs, outputs=[locs, scales, logits])
  return model

K = 20  # number of mixture components
features = ...  # data features

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

You can then train it using gradient descent following any TF 2.0 tutorial.

[Niknafs]

Thanks, Dustin! Can you please verify that the references to Categorical and Normal are from edward2, and not tfp.distributions?

When running the above using Categorical and Normal from edward2, I get the following error:

TypeError: cat must be a Categorical distribution, but saw: RandomVariable("Categorical_1/", shape=(?,), dtype = int32)

Also, do you mind sharing a pointer to one such TF 2.0 tutorial? I am running TF 1.14.0 and TFP 0.7.0.