Skip to content

Instantly share code, notes, and snippets.

@ericjang
Last active January 17, 2018 17:42
Show Gist options
  • Save ericjang/342a0cc08fa6478f4b929ce41b0b1254 to your computer and use it in GitHub Desktop.
Save ericjang/342a0cc08fa6478f4b929ce41b0b1254 to your computer and use it in GitHub Desktop.
d, r = 2, 2
DTYPE = tf.float32
bijectors = []
num_layers = 6
for i in range(num_layers):
with tf.variable_scope('bijector_%d' % i):
V = tf.get_variable('V', [d, r], dtype=DTYPE) # factor loading
shift = tf.get_variable('shift', [d], dtype=DTYPE) # affine shift
L = tf.get_variable('L', [d * (d + 1) / 2],
dtype=DTYPE) # lower triangular
bijectors.append(tfb.Affine(
scale_tril=tfd.fill_triangular(L),
scale_perturb_factor=V,
shift=shift,
))
alpha = tf.abs(tf.get_variable('alpha', [], dtype=DTYPE)) + .01
bijectors.append(LeakyReLU(alpha=alpha))
# Last layer is affine. Note that tfb.Chain takes a list of bijectors in the *reverse* order
# that they are applied.
mlp_bijector = tfb.Chain(
list(reversed(bijectors[:-1])), name='2d_mlp_bijector')
dist = tfd.TransformedDistribution(
distribution=base_dist,
bijector=mlp_bijector
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment