Distributions of high dimensional parameters in Tensorboard

Dear Edward’s users.
I do not understand something using Tensorboard: how can I plot distributions for two dimensional parameters?

For a simple two Gaussian Mixture model (code below) I would expected two means in 2D (one for each component), but when I see the Distributions figure using Tensorboard I only see one mean (in one dimension?) called, in my case, qmu_loc/0 … but the parameter to optimize in qmu distribution is

loc=tf.Variable(tf.random_normal([K,D]),name="qmu_loc),

i.e., shape [2,2]. I don’t understand what Tensorboard does with the dimensions of the parameters.

#!/usr/bin/env python
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import edward as ed
import numpy as np
from scipy.stats import norm
import time
import tensorflow as tf

import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.cm as cm
plt.style.use('ggplot')

from edward.models import \
    Categorical, Empirical, InverseGamma, Normal, Gamma, Beta, Dirichlet, MultivariateNormalDiag, ParamMixture, Mixture

def build_toy_dataset(N):
  pi = np.array([0.3, 0.7])
  mus = [[5, 5], [-1, -1]]
  stds = [[1, 1], [0.1, 0.1]]
  x = np.zeros((N, 2))
  for n in range(N):
    k = np.argmax(np.random.multinomial(1, pi))
    x[n, :] = np.random.multivariate_normal(mus[k], np.diag(stds[k]))
  return x

# dataset

N = 500  # number of data points
x_data = build_toy_dataset(N)
N, D = x_data.shape


K = 2  # number of components to test

with tf.name_scope("model"):
  pi = Dirichlet(concentration=tf.ones(K), name="pi")
  mu = Normal(loc=[[5.0, 5.0], [-1.0, -1.0]], scale=tf.ones([K,D]), name="mu")
  sigmasq = InverseGamma(tf.ones(D), tf.ones(D), sample_shape=K, name="sigmasq")
  z = Categorical(probs=pi, sample_shape=N, name="z")
  components = [MultivariateNormalDiag(loc=mu[k], scale_diag=tf.sqrt(sigmasq[k])) for k in range(K)]
  x = Mixture(cat=z, components=components, sample_shape=N, name="mixture")

with tf.name_scope("posterior"):
  qpi = Dirichlet(tf.nn.softplus(tf.Variable(tf.random_normal([K]),name="qpi")))
  qmu = MultivariateNormalDiag(loc=tf.Variable(tf.random_normal([K,D]),name="qmu_loc"), scale_diag=tf.sqrt([[10.0, 10.0],[10.0, 10.0]]), name="qmu/posterior")
  qsigmasq = InverseGamma(concentration=tf.nn.softplus(tf.Variable(tf.random_normal([K,D]),name="concentration")), rate=tf.ones([K,D]), name="qsigmasq/posterior")
  
inference = ed.KLqp({pi: qpi, mu: qmu, sigmasq: qsigmasq}, data={x: x_data})

inference.run(n_iter=5000, n_samples=30, logdir='log2DMH')

I only see one mean (in one dimension?) called, in my case, qmu_loc/0

qmu_loc/0 is the (unique) name of the tensor. You can check by calling print(loc.name) (or loc.unique_name).

I don’t understand what Tensorboard does with the dimensions of the parameters.

TensorBoard takes a multi-dimensional array and plots all its elements as it changes across training iteration/time. Take a slice at some point, x = x_0. This line represents the distribution of all elements, where different shades correspond to 1, 2, and 3 standard deviations from the mean.

Thanks again @dustin!

It is also specified at Edward-Tensorboard:

Distributions displays the distribution of each non-scalar TensorFlow variable across iterations. These variables are the free parameters of your model and approximating family.

and

Histograms displays the same information as Distributions but as a 3-D histogram changing aross iteration.