Add rebar model.

CODEOWNERS

@@ -19,6 +19,7 @@ object_detection/* @jch1 @tombstone @derekjchow @jesu9 @dreamdragon
 pcl_rl/* @ofirnachum
 ptn/* @xcyan @arkanath @hellojas @honglaklee
 real_nvp/* @laurent-dinh
+rebar/* @gjtucker
 resnet/* @panyx0718
 skip_thoughts/* @cshallue
 slim/* @sguada @nathansilberman

rebar/README.md

+# REINFORCing Concrete with REBAR
+*Implemention of REBAR (and other closely related methods) as described
+in "REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models" by
+George Tucker, Andriy Mnih, Chris J. Maddison, Dieterich Lawson, Jascha Sohl-Dickstein [https://arxiv.org/abs/1703.07370](https://arxiv.org/abs/1703.07370).*
+
+Learning in models with discrete latent variables is challenging due to high variance gradient estimators. Generally, approaches have relied on control variates to reduce the variance of the REINFORCE estimator. Recent work (Jang et al. 2016; Maddison et al. 2016) has taken a different approach, introducing a continuous relaxation of discrete variables to produce low-variance, but biased, gradient estimates. In this work, we combine the two approaches through a novel control variate that produces low-variance, unbiased gradient estimates. Then, we introduce a novel continuous relaxation and show that the tightness of the relaxation can be adapted online, removing it as a hyperparameter. We show state-of-the-art variance reduction on several benchmark generative modeling tasks, generally leading to faster convergence to a better final log likelihood.
+
+REBAR applied to multilayer sigmoid belief networks is implemented in rebar.py and rebar_train.py provides a training/evaluation setup. As a comparison, we also implemented the following methods:
+* [NVIL](https://arxiv.org/abs/1402.0030)
+* [MuProp](https://arxiv.org/abs/1511.05176)
+* [Gumbel-Softmax](https://arxiv.org/abs/1611.01144)
+
+The code is not optimized and some computation is repeated for ease of
+implementation. We hope that this code will be a useful starting point for future research in this area.
+
+## Quick Start:
+
+Requirements:
+* TensorFlow (see tensorflow.org for how to install)
+* MNIST dataset
+* Omniglot dataset
+
+First download datasets, by selecting URLs to download the data from. Then
+fill in the download_data.py script like so:
+
+```
+MNIST_URL = 'http://yann.lecun.com/exdb/mnist'
+MNIST_BINARIZED_URL = 'http://www.cs.toronto.edu/~larocheh/public/datasets/binarized_mnist'
+OMNIGLOT_URL = 'https://github.com/yburda/iwae/raw/master/datasets/OMNIGLOT/chardata.mat'
+```
+
+Then run the script to download the data:
+
+```
+python download_data.py
+```
+
+Then run the training script:
+
+```
+python rebar_train.py --hparams="model=SBNDynamicRebar,learning_rate=0.0003,n_layer=2,task=sbn"
+```
+
+and you should see something like:
+
+```
+Step 2084: [-231.026474      0.3711713     1.            1.06934261    1.07023323
+    1.02173257    1.02171052    1.            1.            1.            1.        ]
+-3.6465678215
+Step 4168: [-156.86795044    0.3097114     1.            1.03964758    1.03936625
+    1.02627242    1.02629256    1.            1.            1.            1.        ]
+-4.42727231979
+Step 6252: [-143.4650116     0.26153237    1.            1.03633797    1.03600132
+    1.02639604    1.02639794    1.            1.            1.            1.        ]
+-4.85577583313
+Step 8336: [-137.65275574    0.22313026    1.            1.03467286    1.03428006
+    1.02336085    1.02335203    0.99999988    1.            0.99999988
+    1.        ]
+-4.95563364029
+```
+
+The first number in the list is the log likelihood lower bound and the number
+after the list is the log of the variance of the gradient estimator. The rest of
+the numbers are for debugging.
+
+We can also compare the variance between methods:
+
+```
+python rebar_train.py \
+  --hparams="model=SBNTrackGradVariances,learning_rate=0.0003,n_layer=2,task=omni"
+```
+
+and you should see something like:
+
+```
+Step 959: [ -2.60478699e+02   3.84281784e-01   6.31126612e-02   3.27319391e-02
+   6.13379292e-03   1.98278503e-04   1.96425783e-04   8.83973844e-04
+   8.70995224e-04             -inf]
+('DynamicREBAR', -3.725339889526367)
+('MuProp', -0.033569782972335815)
+('NVIL', 2.7640280723571777)
+('REBAR', -3.539274215698242)
+('SimpleMuProp', -0.040744658559560776)
+Step 1918: [ -2.06948471e+02   3.35904926e-01   5.20901568e-03   7.81541676e-05
+   2.06885766e-03   1.08521657e-04   1.07351625e-04   2.30646547e-04
+   2.26554010e-04  -8.22885323e+00]
+('DynamicREBAR', -3.864381790161133)
+('MuProp', -0.7183765172958374)
+('NVIL', 2.266523599624634)
+('REBAR', -3.662022113800049)
+('SimpleMuProp', -0.7071359157562256)
+```
+where the tuples show the log of the variance of the gradient estimators.
+
+The training script has a number of hyperparameter configuration flags:
+* task (sbn): one of {sbn, sp, omni} which correspond to MNIST generative
+  modeling, structured prediction on MNIST, and Omniglot generative modeling,
+  respectively
+* model (SBNGumbel) : one of {SBN, SBNNVIL, SBNMuProp, SBNSimpleMuProp,
+  SBNRebar, SBNDynamicRebar, SBNGumbel SBNTrackGradVariances}. DynamicRebar automatically
+  adjusts the temperature, whereas Rebar and Gumbel-Softmax require tuning the
+  temperature. The ones named after
+  methods uses that method to estimate the gradients (SBN refers to
+  REINFORCE). SBNTrackGradVariances runs multiple methods and follows a single
+  optimization trajectory.
+* n_hidden (200): number of hidden nodes per layer
+* n_layer (1): number of layers in the model
+* nonlinear (false): if true use 2 x tanh layers between each stochastic layer,
+  otherwise use a linear layer
+* learning_rate (0.001): learning rate
+* temperature (0.5): temperature hyperparameter (for DynamicRebar, this is the initial
+  value of the temperature)
+* n_samples (1): number of samples used to compute the gradient estimator (for the
+  experiments in the paper, set to 1)
+* batch_size (24): batch size
+* muprop_relaxation (true): if true use the new relaxation described in the paper,
+  otherwise use the Concrete/Gumbel softmax relaxation
+* dynamic_b (false): if true dynamically binarize the training set. This
+  increases the effective training dataset size and reduces overfitting, though
+  it is not a standard dataset
+
+Maintained by George Tucker (gjt@google.com, github user: gjtucker).

rebar/config.py

+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Configuration variables."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+DATA_DIR = 'data'
+MNIST_BINARIZED = 'mnist_salakhutdinov_07-19-2017.pkl'
+MNIST_FLOAT = 'mnist_train_xs_07-19-2017.npy'
+OMNIGLOT = 'omniglot_07-19-2017.mat'

rebar/datasets.py

+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Library of datasets for REBAR."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import random
+import os
+import scipy.io
+import numpy as np
+import cPickle as pickle
+import tensorflow as tf
+import config
+gfile = tf.gfile
+
+
+def load_data(hparams):
+  # Load data
+  if hparams.task in ['sbn', 'sp']:
+    reader = read_MNIST
+  elif hparams.task == 'omni':
+    reader = read_omniglot
+  x_train, x_valid, x_test = reader(binarize=not hparams.dynamic_b)
+
+  return x_train, x_valid, x_test
+
+def read_MNIST(binarize=False):
+  """Reads in MNIST images.
+
+  Args:
+    binarize: whether to use the fixed binarization
+
+  Returns:
+    x_train: 50k training images
+    x_valid: 10k validation images
+    x_test: 10k test images
+
+  """
+  with gfile.FastGFile(os.path.join(config.DATA_DIR, config.MNIST_BINARIZED), 'r') as f:
+    (x_train, _), (x_valid, _), (x_test, _) = pickle.load(f)
+
+  if not binarize:
+    with gfile.FastGFile(os.path.join(config.DATA_DIR, config.MNIST_FLOAT), 'r') as f:
+      x_train = np.load(f).reshape(-1, 784)
+
+  return x_train, x_valid, x_test
+
+def read_omniglot(binarize=False):
+  """Reads in Omniglot images.
+
+  Args:
+    binarize: whether to use the fixed binarization
+
+  Returns:
+    x_train: training images
+    x_valid: validation images
+    x_test: test images
+
+  """
+  n_validation=1345
+
+  def reshape_data(data):
+    return data.reshape((-1, 28, 28)).reshape((-1, 28*28), order='fortran')
+
+  omni_raw = scipy.io.loadmat(os.path.join(config.DATA_DIR, config.OMNIGLOT))
+
+  train_data = reshape_data(omni_raw['data'].T.astype('float32'))
+  test_data = reshape_data(omni_raw['testdata'].T.astype('float32'))
+
+  # Binarize the data with a fixed seed
+  if binarize:
+    np.random.seed(5)
+    train_data = (np.random.rand(*train_data.shape) < train_data).astype(float)
+    test_data = (np.random.rand(*test_data.shape) < test_data).astype(float)
+
+  shuffle_seed = 123
+  permutation = np.random.RandomState(seed=shuffle_seed).permutation(train_data.shape[0])
+  train_data = train_data[permutation]
+
+  x_train = train_data[:-n_validation]
+  x_valid = train_data[-n_validation:]
+  x_test = test_data
+
+  return x_train, x_valid, x_test
+

rebar/download_data.py

+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Download MNIST, Omniglot datasets for Rebar."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import urllib
+import gzip
+import os
+import config
+import struct
+import numpy as np
+import cPickle as pickle
+import datasets
+
+MNIST_URL = 'see README'
+MNIST_BINARIZED_URL = 'see README'
+OMNIGLOT_URL = 'see README'
+
+MNIST_FLOAT_TRAIN = 'train-images-idx3-ubyte'
+
+
+def load_mnist_float(local_filename):
+  with open(local_filename, 'rb') as f:
+    f.seek(4)
+    nimages, rows, cols = struct.unpack('>iii', f.read(12))
+    dim = rows*cols
+
+    images = np.fromfile(f, dtype=np.dtype(np.ubyte))
+    images = (images/255.0).astype('float32').reshape((nimages, dim))
+
+  return images
+
+if __name__ == '__main__':
+  if not os.path.exists(config.DATA_DIR):
+    os.makedirs(config.DATA_DIR)
+
+  # Get MNIST and convert to npy file
+  local_filename = os.path.join(config.DATA_DIR, MNIST_FLOAT_TRAIN)
+  if not os.path.exists(local_filename):
+    urllib.urlretrieve("%s/%s.gz" % (MNIST_URL, MNIST_FLOAT_TRAIN), local_filename+'.gz')
+    with gzip.open(local_filename+'.gz', 'rb') as f:
+      file_content = f.read()
+    with open(local_filename, 'wb') as f:
+      f.write(file_content)
+    os.remove(local_filename+'.gz')
+
+  mnist_float_train = load_mnist_float(local_filename)[:-10000]
+  # save in a nice format
+  np.save(os.path.join(config.DATA_DIR, config.MNIST_FLOAT), mnist_float_train)
+
+  # Get binarized MNIST
+  splits = ['train', 'valid', 'test']
+  mnist_binarized = []
+  for split in splits:
+    filename = 'binarized_mnist_%s.amat' % split
+    url = '%s/binarized_mnist_%s.amat' % (MNIST_BINARIZED_URL, split)
+    local_filename = os.path.join(config.DATA_DIR, filename)
+    if not os.path.exists(local_filename):
+      urllib.urlretrieve(url, local_filename)
+
+    with open(local_filename, 'rb') as f:
+      mnist_binarized.append((np.array([map(int, line.split()) for line in f.readlines()]).astype('float32'), None))
+
+  # save in a nice format
+  with open(os.path.join(config.DATA_DIR, config.MNIST_BINARIZED), 'w') as out:
+    pickle.dump(mnist_binarized, out)
+
+  # Get Omniglot
+  local_filename = os.path.join(config.DATA_DIR, config.OMNIGLOT)
+  if not os.path.exists(local_filename):
+    urllib.urlretrieve(OMNIGLOT_URL,
+                       local_filename)
+

rebar/logger.py

+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Logger for REBAR"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+class Logger:
+  def __init__(self):
+    pass
+
+  def log(self, key, value):
+    pass
+
+  def flush(self):
+    pass

rebar/rebar_train.py

+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+import random
+import sys
+import os
+
+import numpy as np
+import tensorflow as tf
+
+import rebar
+import datasets
+import logger as L
+gfile = tf.gfile
+
+tf.app.flags.DEFINE_string("working_dir", "/tmp/rebar",
+                           """Directory where to save data, write logs, etc.""")
+tf.app.flags.DEFINE_string('hparams', '',
+                           '''Comma separated list of name=value pairs.''')
+tf.app.flags.DEFINE_integer('eval_freq', 20,
+                           '''How often to run the evaluation step.''')
+FLAGS = tf.flags.FLAGS
+
+def manual_scalar_summary(name, value):
+  value = tf.Summary.Value(tag=name, simple_value=value)
+  summary_str = tf.Summary(value=[value])
+  return summary_str
+
+def eval(sbn, eval_xs, n_samples=100, batch_size=5):
+  n = eval_xs.shape[0]
+  i = 0
+  res = []
+  while i < n:
+    batch_xs = eval_xs[i:min(i+batch_size, n)]
+    res.append(sbn.partial_eval(batch_xs, n_samples))
+    i += batch_size
+  res = np.mean(res, axis=0)
+  return res
+
+def train(sbn, train_xs, valid_xs, test_xs, training_steps, debug=False):
+  hparams = sorted(sbn.hparams.values().items())
+  hparams = (map(str, x) for x in hparams)
+  hparams = ('_'.join(x) for x in hparams)
+  hparams_str = '.'.join(hparams)
+
+  logger = L.Logger()
+
+  # Create the experiment name from the hparams
+  experiment_name = ([str(sbn.hparams.n_hidden) for i in xrange(sbn.hparams.n_layer)] +
+                     [str(sbn.hparams.n_input)])
+  if sbn.hparams.nonlinear:
+    experiment_name = '~'.join(experiment_name)
+  else:
+    experiment_name = '-'.join(experiment_name)
+  experiment_name = 'SBN_%s' % experiment_name
+  rowkey = {'experiment': experiment_name,
+            'model': hparams_str}
+
+  # Create summary writer
+  summ_dir = os.path.join(FLAGS.working_dir, hparams_str)
+  summary_writer = tf.summary.FileWriter(
+      summ_dir, flush_secs=15, max_queue=100)
+
+  sv = tf.train.Supervisor(logdir=os.path.join(
+      FLAGS.working_dir, hparams_str),
+                     save_summaries_secs=0,
+                     save_model_secs=1200,
+                     summary_op=None,
+                     recovery_wait_secs=30,
+                     global_step=sbn.global_step)
+  with sv.managed_session() as sess:
+    # Dump hparams to file
+    with gfile.Open(os.path.join(FLAGS.working_dir,
+                                 hparams_str,
+                                 'hparams.json'),
+                    'w') as out:
+      json.dump(sbn.hparams.values(), out)
+
+    sbn.initialize(sess)
+    batch_size = sbn.hparams.batch_size
+    scores = []
+    n = train_xs.shape[0]
+    index = range(n)
+
+    while not sv.should_stop():
+      lHats = []
+      grad_variances = []
+      temperatures = []
+      random.shuffle(index)
+      i = 0
+      while i < n:
+        batch_index = index[i:min(i+batch_size, n)]
+        batch_xs = train_xs[batch_index, :]
+
+        if sbn.hparams.dynamic_b:
+          # Dynamically binarize the batch data
+          batch_xs = (np.random.rand(*batch_xs.shape) < batch_xs).astype(float)
+
+        lHat, grad_variance, step, temperature = sbn.partial_fit(batch_xs,
+                                                    sbn.hparams.n_samples)
+        if debug:
+          print(i, lHat)
+          if i > 100:
+            return
+        lHats.append(lHat)
+        grad_variances.append(grad_variance)
+        temperatures.append(temperature)
+        i += batch_size
+
+      grad_variances = np.log(np.mean(grad_variances, axis=0)).tolist()
+      summary_strings = []
+      if isinstance(grad_variances, list):
+        grad_variances = dict(zip([k for (k, v) in sbn.losses], map(float, grad_variances)))
+        rowkey['step'] = step
+        logger.log(rowkey, {'step': step,
+                             'train': np.mean(lHats, axis=0)[0],
+                             'grad_variances': grad_variances,
+                             'temperature': np.mean(temperatures), })
+        grad_variances = '\n'.join(map(str, sorted(grad_variances.iteritems())))
+      else:
+        rowkey['step'] = step
+        logger.log(rowkey, {'step': step,
+                             'train': np.mean(lHats, axis=0)[0],
+                             'grad_variance': grad_variances,
+                             'temperature': np.mean(temperatures), })
+        summary_strings.append(manual_scalar_summary("log grad variance", grad_variances))
+
+      print('Step %d: %s\n%s' % (step, str(np.mean(lHats, axis=0)), str(grad_variances)))
+
+      # Every few epochs compute test and validation scores
+      epoch = int(step / (train_xs.shape[0] / sbn.hparams.batch_size))
+      if epoch % FLAGS.eval_freq == 0:
+        valid_res = eval(sbn, valid_xs)
+        test_res= eval(sbn, test_xs)
+
+        print('\nValid %d: %s' % (step, str(valid_res)))
+        print('Test %d: %s\n' % (step, str(test_res)))
+        logger.log(rowkey, {'step': step,
+                             'valid': valid_res[0],
+                             'test': test_res[0]})
+        logger.flush()  # Flush infrequently
+
+      # Create summaries
+      summary_strings.extend([
+        manual_scalar_summary("Train ELBO", np.mean(lHats, axis=0)[0]),
+        manual_scalar_summary("Temperature", np.mean(temperatures)),
+      ])
+      for summ_str in summary_strings:
+        summary_writer.add_summary(summ_str, global_step=step)
+      summary_writer.flush()
+
+      sys.stdout.flush()
+      scores.append(np.mean(lHats, axis=0))
+
+      if step > training_steps:
+        break
+
+    return scores
+
+
+def main():
+  # Parse hyperparams
+  hparams = rebar.default_hparams
+  hparams.parse(FLAGS.hparams)
+  print(hparams.values())
+
+  train_xs, valid_xs, test_xs = datasets.load_data(hparams)
+  mean_xs = np.mean(train_xs, axis=0)  # Compute mean centering on training
+
+  training_steps = 2000000
+  model = getattr(rebar, hparams.model)
+  sbn = model(hparams, mean_xs=mean_xs)
+
+  scores = train(sbn, train_xs, valid_xs, test_xs,
+                 training_steps=training_steps, debug=False)
+
+if __name__ == '__main__':
+  main()

rebar/utils.py

+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Basic data management and plotting utilities."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import cPickle as pickle
+import getpass
+import numpy as np
+import gc
+import tensorflow as tf
+
+#
+# Python utlities
+#
+def exp_moving_average(x, alpha=0.9):
+  res = []
+  mu = 0
+  alpha_factor = 1
+  for x_i in x:
+    mu += (1 - alpha)*(x_i - mu)
+    alpha_factor *= alpha
+    res.append(mu/(1 - alpha_factor))
+
+  return np.array(res)
+
+def sanitize(s):
+  return s.replace('.', '_')
+
+#
+# Tensorflow utilities
+#
+def softplus(x):
+  '''
+  Let m = max(0, x), then,
+
+  sofplus(x) = log(1 + e(x)) = log(e(0) + e(x)) = log(e(m)(e(-m) + e(x-m)))
+             = m + log(e(-m) + e(x - m))
+
+  The term inside of the log is guaranteed to be between 1 and 2.
+  '''
+  m = tf.maximum(tf.zeros_like(x), x)
+  return m + tf.log(tf.exp(-m) + tf.exp(x - m))
+
+def safe_log_prob(x, eps=1e-8):
+  return tf.log(tf.clip_by_value(x, eps, 1.0))
+
+def rms(x):
+  return tf.sqrt(tf.reduce_mean(tf.square(x)))
+
+def center(x):
+  mu = (tf.reduce_sum(x) - x)/tf.to_float(tf.shape(x)[0] - 1)
+  return x - mu
+
+def vectorize(grads_and_vars, set_none_to_zero=False, skip_none=False):
+  if set_none_to_zero:
+    return tf.concat([tf.reshape(g, [-1]) if g is not None else
+                         tf.reshape(tf.zeros_like(v), [-1]) for g, v in grads_and_vars], 0)
+  elif skip_none:
+    return tf.concat([tf.reshape(g, [-1]) for g, v in grads_and_vars if g is not None], 0)
+  else:
+    return tf.concat([tf.reshape(g, [-1]) for g, v in grads_and_vars], 0)
+
+def add_grads_and_vars(a, b):
+  '''Add grads_and_vars from two calls to tf.compute_gradients.'''
+  res = []
+  for (g_a, v_a), (g_b, v_b) in zip(a, b):
+    assert v_a == v_b
+    if g_a is None:
+      res.append((g_b, v_b))
+    elif g_b is None:
+      res.append((g_a, v_a))
+    else:
+      res.append((g_a + g_b, v_a))
+  return res
+
+def binary_log_likelihood(y, log_y_hat):
+  """Computes binary log likelihood.
+
+  Args:
+    y: observed data
+    log_y_hat: parameters of the binary variables
+
+  Returns:
+    log_likelihood
+  """
+  return tf.reduce_sum(y*(-softplus(-log_y_hat)) +
+                       (1 - y)*(-log_y_hat-softplus(-log_y_hat)),
+                       1)
+
+def cov(a, b):
+  """Compute the sample covariance between two vectors."""
+  mu_a = tf.reduce_mean(a)
+  mu_b = tf.reduce_mean(b)
+  n = tf.to_float(tf.shape(a)[0])
+
+  return tf.reduce_sum((a - mu_a)*(b - mu_b))/(n - 1.0)
+
+def corr(a, b):
+  return cov(a, b)*tf.rsqrt(cov(a, a))*tf.rsqrt(cov(b, b))
+
+def logSumExp(t, axis=0, keep_dims = False):
+  '''Computes the log(sum(exp(t))) numerically stabily.
+
+  Args:
+    t: input tensor
+    axis: which axis to sum over
+    keep_dims: whether to keep the dim or not
+
+  Returns:
+    tensor with result
+
+  '''
+  m = tf.reduce_max(t, [axis])
+  res = m + tf.log(tf.reduce_sum(tf.exp(t - tf.expand_dims(m, axis)), [axis]))
+
+  if keep_dims:
+    return tf.expand_dims(res, axis)
+  else:
+    return res
+
+if __name__ == '__main__':
+  app.run()