MDN Implementation using Edward2 about edward2 HOT 2 OPEN

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))