```

```

[](https://colab.research.google.com/github/m2lschool/tutorials2021/blob/main/3_generative/VAE_Tutorial_Start.ipynb)

Contact: {fviola@google.com, marco.ciccone@me.com}

## Contents
* Enabling TPUs in colab
* Handling nested data structures using tree utilities in JAX
* Distributing computation over multiple devices using *pmap*
* Amortized variational inference (VAEs)
  * Training VAEs optimizing ELBO
  * Training $\beta$-VAEs
  * Training VAEs using constraint optimization (GECO)

# Set up your environment!

```
#@title Download and install all the missing packages required for this tutorial { display-mode: "code" }
! pip install ipdb -q
! pip install chex -q
! pip install optax -q
! pip install dm_haiku -q
! pip install tfp-nightly[jax] -q
! pip install tf-nightly -q
! pip install livelossplot -q
# ! pip install gast==0.3.3 -q
# ! pip install grpcio==1.32.0 -q
# ! pip install h5py==2.10.0 -q

print("All packages installed!")
```

    ERROR: tensorflow 2.4.0 has requirement gast==0.3.3, but you'll have gast 0.4.0 which is incompatible.
    ERROR: tensorflow 2.4.0 has requirement grpcio~=1.32.0, but you'll have grpcio 1.34.0 which is incompatible.
    ERROR: tensorflow 2.4.0 has requirement h5py~=2.10.0, but you'll have h5py 3.1.0 which is incompatible.
    All packages installed!

```
# @title Imports
import inspect
import os

import chex
import dill
import functools
import haiku as hk
import jax
import jax.numpy as jnp

import numpy as np
import optax as tx
import requests
from pprint import pprint

import seaborn as sns
from matplotlib import pyplot as plt
from livelossplot import PlotLosses
from livelossplot.outputs import MatplotlibPlot

import tensorflow as tf
import tensorflow_probability
from tensorflow_probability.substrates import jax as tfp

sns.set(rc={"lines.linewidth": 2.8}, font_scale=2)
sns.set_style("whitegrid")

# Returns the code of the python implementation of a given funciton as a string.
get_code_as_string = lambda fn: dill.source.getsource(fn.__code__)
```

```
# ------------------ #
# Enable TPU support #
# ------------------ #

# This cell execution might take a while! don't worry :)
# Don't forget to select a TPU or GPU runtime environment in
# Runtime -> Change runtime type
try:
  if 'TPU_DRIVER_MODE' not in globals():
    url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver_nightly'
    resp = requests.post(url)
    TPU_DRIVER_MODE = 1

  # The following is required to use TPU Driver as JAX's backend.
  from jax.config import config
  config.FLAGS.jax_xla_backend = "tpu_driver"
  config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
except:
  print('TPUs not found. Enable a TPU runtime going to: '
        '"Runtime -> Change runtime type"')
devices = jax.devices()
print("Available devices:", devices)
# Should print something like
# Available devices: [TpuDevice(id=0, host_id=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, host_id=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, host_id=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, host_id=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, host_id=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, host_id=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, host_id=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, host_id=0, coords=(1,1,0), core_on_chip=1)]
```

    Available devices: [TpuDevice(id=0, host_id=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, host_id=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, host_id=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, host_id=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, host_id=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, host_id=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, host_id=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, host_id=0, coords=(1,1,0), core_on_chip=1)]

```
devices
```

    [TpuDevice(id=0, host_id=0, coords=(0,0,0), core_on_chip=0),
     TpuDevice(id=1, host_id=0, coords=(0,0,0), core_on_chip=1),
     TpuDevice(id=2, host_id=0, coords=(1,0,0), core_on_chip=0),
     TpuDevice(id=3, host_id=0, coords=(1,0,0), core_on_chip=1),
     TpuDevice(id=4, host_id=0, coords=(0,1,0), core_on_chip=0),
     TpuDevice(id=5, host_id=0, coords=(0,1,0), core_on_chip=1),
     TpuDevice(id=6, host_id=0, coords=(1,1,0), core_on_chip=0),
     TpuDevice(id=7, host_id=0, coords=(1,1,0), core_on_chip=1)]

# Warmup: bits and bobs

Let's take a minute to look at some JAX functionality we will be using in this tutorial.

## How to work with nested data with `jax.tree_utils`

It is fairly common to structure data and parameters into nested formats, for example (nested) dictionaries, namedtuples, lists, dataclasses and variants thereof. For example, we might want to cleanly scope and group our model's variables by their component name into some dictionary-like format, like in Haiku, or in a Reinforcement Learning setting we might want to explicitely label the components of a rollout.

Typically in our code we don't want to assume a priori any strucuture of the containers we are manipulating, and prefer code that can transparently handle arbitrary nested structures. Luckily JAX can natively do this for us, and we only need to familiarize ourselves with its `jax.tree_util` package, and make sure that our custom objects are registered with it (not to worry, we have libraries that do this for us!).

You can find out more in the JAX `tree_util` package [documentation](https://jax.readthedocs.io/en/latest/jax.tree_util.html).

```
from collections import namedtuple
data_container = namedtuple('data_box', 'component_a component_b')
data = dict(
      a=jnp.ones(shape=()),
      b= [
          jnp.ones(shape=()),
          data_container(jnp.ones(shape=()), jnp.ones(shape=()))],
    c=(jnp.ones(shape=()), jnp.ones(shape=())))
print('Structured data\n', data)
```

    Structured data
     {'a': DeviceArray(1., dtype=float32), 'b': [DeviceArray(1., dtype=float32), data_box(component_a=DeviceArray(1., dtype=float32), component_b=DeviceArray(1., dtype=float32))], 'c': (DeviceArray(1., dtype=float32), DeviceArray(1., dtype=float32))}

```
data
```

    {'a': DeviceArray(1., dtype=float32),
     'b': [DeviceArray(1., dtype=float32),
      data_box(component_a=DeviceArray(1., dtype=float32), component_b=DeviceArray(1., dtype=float32))],
     'c': (DeviceArray(1., dtype=float32), DeviceArray(1., dtype=float32))}

```
# We can use `jax.tree_map` to apply the same function to all the tensors
# contained in a nested data structure.
fn = lambda x: x * 2
output = jax.tree_map(fn, data)
print('Structure data, after {}'.format(get_code_as_string(fn)), output)
```

    Structure data, after fn = lambda x: x * 2
     {'a': DeviceArray(2., dtype=float32), 'b': [DeviceArray(2., dtype=float32), data_box(component_a=DeviceArray(2., dtype=float32), component_b=DeviceArray(2., dtype=float32))], 'c': (DeviceArray(2., dtype=float32), DeviceArray(2., dtype=float32))}

```
data = dict(
      a = jnp.array([1]),
      b = [
          jnp.array([2]),
          data_container(jnp.array([3]), jnp.array([4]))
      ],
      c = (jnp.array([6]), jnp.array([5]))
)
print('Structured data\n', data)
```

    Structured data
     {'a': DeviceArray([1], dtype=int32), 'b': [DeviceArray([2], dtype=int32), data_box(component_a=DeviceArray([3], dtype=int32), component_b=DeviceArray([4], dtype=int32))], 'c': (DeviceArray([6], dtype=int32), DeviceArray([5], dtype=int32))}

```
# We can also call functions with multiple structured inputs, for example
# parameters and gradients in an update step.
fn = lambda x, y, delta=0.1: x + delta * y
output = jax.tree_multimap(fn, data, data)
print('Structure data, after {}'.format(get_code_as_string(fn)), output)
```

    Structure data, after fn = lambda x, y, delta=0.1: x + delta * y
     {'a': DeviceArray([1.1], dtype=float32), 'b': [DeviceArray([2.2], dtype=float32), data_box(component_a=DeviceArray([3.3], dtype=float32), component_b=DeviceArray([4.4], dtype=float32))], 'c': (DeviceArray([6.6], dtype=float32), DeviceArray([5.5], dtype=float32))}

```

```

```
# We can also 'flatten' the data to get a list of all the tensors contained in
# the nested data structure.
entries = jax.tree_leaves(data)
print('Tree leaves\n', entries)
```

    Tree leaves
     [DeviceArray([1], dtype=int32), DeviceArray([2], dtype=int32), DeviceArray([3], dtype=int32), DeviceArray([4], dtype=int32), DeviceArray([6], dtype=int32), DeviceArray([5], dtype=int32)]

```

# We can 'unflatten' flattened data to get back the original strucure if we keep
# around the original structure defintion.
entries, tree_def = jax.tree_flatten(data)  # tree_def capture the structure
data = jax.tree_unflatten(tree_def, entries)
print('Flattened data\n', entries)
print('Unflattened data\n', data)
```

    Flattened data
     [DeviceArray([1], dtype=int32), DeviceArray([2], dtype=int32), DeviceArray([3], dtype=int32), DeviceArray([4], dtype=int32), DeviceArray([6], dtype=int32), DeviceArray([5], dtype=int32)]
    Unflattened data
     {'a': DeviceArray([1], dtype=int32), 'b': [DeviceArray([2], dtype=int32), data_box(component_a=DeviceArray([3], dtype=int32), component_b=DeviceArray([4], dtype=int32))], 'c': (DeviceArray([6], dtype=int32), DeviceArray([5], dtype=int32))}

```
tree_def
```

    PyTreeDef(dict[['a', 'b', 'c']], [*,PyTreeDef(list, [*,PyTreeDef(namedtuple[<class '__main__.data_box'>], [*,*])]),PyTreeDef(tuple, [*,*])])

```
# ------------------------ #
# WARMUP OPTIONAL EXERCISE #
# ------------------------ #

# You have some batched data, structured in an unknow way and you want
# to recover a list of unbatched data, structured the same way.
# Write the code to do that using jax.tree_utils

input_data = dict(
    a=jnp.arange(3),
    b=[jnp.arange(3),
       data_container(jnp.arange(3), jnp.arange(3))],
    c=(jnp.arange(3), jnp.arange(3)))

print('Input data')
pprint(input_data)

```

    Input data
    {'a': DeviceArray([0, 1, 2], dtype=int32),
     'b': [DeviceArray([0, 1, 2], dtype=int32),
           data_box(component_a=DeviceArray([0, 1, 2], dtype=int32), component_b=DeviceArray([0, 1, 2], dtype=int32))],
     'c': (DeviceArray([0, 1, 2], dtype=int32),
           DeviceArray([0, 1, 2], dtype=int32))}

```
flattened_data, tree_def = jax.tree_flatten(input_data)
```

```
split_data = list(zip(*flattened_data))
```

```
jax.tree_unflatten(tree_def, split_data[0])
```

    {'a': 0, 'b': [0, data_box(component_a=0, component_b=0)], 'c': (0, 0)}

```
jnp.array(flattened_data)[:,1]
```

    DeviceArray([1, 1, 1, 1, 1, 1], dtype=int32)

```
unbatched_data = [jax.tree_unflatten(tree_def, entries) for entries in split_data]
unbatched_data
```

    []

```
flattened_data, tree_def = jax.tree_flatten(input_data)
flattened_data
```

    [DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32)]

```
flattened_data = jnp.array(flattened_data)
flattened_data[:,1].reshape(1,-1)
```

    DeviceArray([[1, 1, 1, 1, 1, 1]], dtype=int32)

```
unbatched_data = [jax.tree_unflatten(tree_def, flattened_data[:,i]) for i in range(len(flattened_data))]
unbatched_data
```

```
flattened_data, tree_def = jax.tree_flatten(input_data)
flattened_data
```

    [DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32),
     DeviceArray([0, 1, 2], dtype=int32)]

```

def unbatch(data):
  flattened_data, tree_def = jax.tree_flatten(data)
  # ADD CODE BELOW
  # -----------------------
  split_data = zip(*flattened_data)
  unbatched_data = [jax.tree_unflatten(tree_def, entries) for entries in split_data]
  # -----------------------
  return unbatched_data

print('\nSplit data')
pprint(unbatch(input_data))

# Expected output:
# Input data
# {'a': DeviceArray([0, 1, 2], dtype=int32),
#  'b': [DeviceArray([0, 1, 2], dtype=int32),
#        data_box(component_a=DeviceArray([0, 1, 2], dtype=int32), component_b=DeviceArray([0, 1, 2], dtype=int32))],
#  'c': (DeviceArray([0, 1, 2], dtype=int32),
#        DeviceArray([0, 1, 2], dtype=int32))}

# Split data
# [{'a': DeviceArray([0], dtype=int32),
#   'b': [DeviceArray([0], dtype=int32),
#         data_box(component_a=DeviceArray([0], dtype=int32), component_b=DeviceArray([0], dtype=int32))],
#   'c': (DeviceArray([0], dtype=int32), DeviceArray([0], dtype=int32))},
#  {'a': DeviceArray([1], dtype=int32),
#   'b': [DeviceArray([1], dtype=int32),
#         data_box(component_a=DeviceArray([1], dtype=int32), component_b=DeviceArray([1], dtype=int32))],
#   'c': (DeviceArray([1], dtype=int32), DeviceArray([1], dtype=int32))},
#  {'a': DeviceArray([2], dtype=int32),
#   'b': [DeviceArray([2], dtype=int32),
#         data_box(component_a=DeviceArray([2], dtype=int32), component_b=DeviceArray([2], dtype=int32))],
#   'c': (DeviceArray([2], dtype=int32), DeviceArray([2], dtype=int32))}]
```

    Split data
    [{'a': 0, 'b': [0, data_box(component_a=0, component_b=0)], 'c': (0, 0)},
     {'a': 1, 'b': [1, data_box(component_a=1, component_b=1)], 'c': (1, 1)},
     {'a': 2, 'b': [2, data_box(component_a=2, component_b=2)], 'c': (2, 2)}]

## How to parallelize gradients computation over multiple devices using JAX.

In JAX you can use the `pmap` primitive to parallelize computation over multiple devices; using reduction methods in `jax.lax` you can also have access to values aggregated across devices during the distributed computation, giving you a lot of flexibility over what you can achieve!

For example, you can split a *large* batch over 8 TPU cores, compute partial
gradients over the split batches and average them prior to updating
the model parameters - in parallel - across all devices.

Let's look at some example code of how this can be done.

```
# ------- #
# EXAMPLE #
# ------- #

# First off, let's verify that you that in this colab you should have access to
# multiple TPU cores.
num_dev = jax.local_device_count()
print('Number of devices:\n', num_dev, '\n')

# Let's define a simple data generation function, and an MSE loss for linear
# regression.
def get_data(prng_key, batch_size):
  x_key, noise_key = jax.random.split(prng_key, 2)
  x = jax.random.uniform(x_key, shape=(batch_size,1))
  true_model = lambda x: x * 3 + 2
  noise = lambda key: jax.random.normal(key, shape=(batch_size, 1)) * 0.2
  return jnp.concatenate([x, true_model(x) + noise(noise_key)], axis=1)

def loss_fn(params, data):
  prediction = params[0] * data[:, 0] + params[1]
  target = data[:, 1]
  return jnp.mean((prediction - target) ** 2)

batch_size = 1024
prng_key = jax.random.PRNGKey(0)
data = get_data(prng_key, batch_size)
params = jax.random.normal(prng_key, shape=(2,))

# On a single core we can compute the gradients over a minibatch with:
grad_fn = jax.grad(loss_fn)
gradients = grad_fn(params, data)
# and update the parameters using any update rule, like those you can find in
# the optax package.

# We now want to take advantage of the the multiple core made available to us.
# To do so, we batch computation _over cores_.

# We format the data and batch it twice, over cores and per-device core:
data = get_data(prng_key, batch_size * num_dev)
data = jnp.reshape(data, (num_dev, batch_size) + data.shape[1:])
# Note that the leading dimension is now the number of cores!
print('Data shape:\n', data.shape, '\n')

# We provide each core with a copy of the parameters. Given an instance of the
# params we can use JAX's broadcasting utility to achieve this.
broadcast = lambda params: jnp.broadcast_to(params, (num_dev,) + params.shape)
mapped_params = broadcast(params)
# All the parameters copies are synced at the start of the model fitting; since
# the udpates to the params will be the same, params will stay synced during
# optimization.
print('Params:\n', params, '\n')
print('Mapped params:\n', mapped_params, '\n')

# We then pmap the gradient fn, averaging the gradients computed across devices
# using the pmap primitive and the jax.lax.pmean function.
def get_averaged_grads(params, data):
  grads = grad_fn(params, data)
  grads = jax.lax.pmean(grads, axis_name='i')
  return grads
get_averaged_grads = jax.pmap(
    get_averaged_grads, axis_name='i', devices=jax.devices())
averaged_grads = get_averaged_grads(mapped_params, data)
print('Averaged_grads:\n', averaged_grads, '\n')

# Note that this is equivalent to mapping the gradient function and manually
# averaging the result!
get_mapped_grads = jax.pmap(grad_fn, axis_name='i', devices=jax.devices())
mapped_grads = get_mapped_grads(mapped_params, data)
print('Mapped_grads:\n', mapped_grads, '\n')
print('Manually averaged_grads:\n', jnp.mean(mapped_grads, axis=0), '\n')

# Of course we can pmap in one step, cleanly, much more complicated logic.
# For example, the whole update step:
optimizer = tx.sgd(1e-1)
def update(params, opt_state, data):
  loss, grads = jax.value_and_grad(loss_fn)(params, data)
  grads = jax.lax.pmean(grads, axis_name='i')
  raw_updates, opt_state = optimizer.update(grads, opt_state)
  params = tx.apply_updates(params, raw_updates)
  return params, opt_state, loss
# Note that pmap also compiles the function to XLA, akin to JIT!
update = jax.pmap(update, axis_name='i', devices=jax.devices())
opt_state = jax.tree_map(broadcast, optimizer.init(params))

losses = []
for _ in range(250):
  prng_key, data_key = jax.random.split(prng_key)
  data = get_data(prng_key, batch_size * num_dev)
  data = jnp.reshape(data, (num_dev, batch_size) + data.shape[1:])
  mapped_params, opt_state, loss = update(mapped_params, opt_state, data)
  losses.append(jnp.mean(loss))

sz = 5; plt.figure(figsize=(4*sz,sz))
plt.subplot(121)
plt.plot(losses)
plt.title('loss')
plt.subplot(122)
data = get_data(prng_key, 100)
plt.scatter(data[:,0], data[:,1], label='data')
plt.plot(
    [0, 1],
    [mapped_params[0, 1], mapped_params[0, 0] + mapped_params[0, 1]],
    'r', label='prediction')
plt.legend();
```

# Get the data

In this tutorial we will use the MNIST and Fashion MNIST datasets (and variations thereof).
We can use TensorFlow data to download the data from the cloud.

```
import tensorflow_datasets as tfds
mnist = tfds.load("mnist")
fashion_mnist = tfds.load("fashion_mnist")
```

    Downloading and preparing dataset mnist/3.0.1 (download: 11.06 MiB, generated: 21.00 MiB, total: 32.06 MiB) to /root/tensorflow_datasets/mnist/3.0.1...

    WARNING:absl:Dataset mnist is hosted on GCS. It will automatically be downloaded to your
    local data directory. If you'd instead prefer to read directly from our public
    GCS bucket (recommended if you're running on GCP), you can instead pass
    `try_gcs=True` to `tfds.load` or set `data_dir=gs://tfds-data/datasets`.

    HBox(children=(FloatProgress(value=0.0, description='Dl Completed...', max=4.0, style=ProgressStyle(descriptio…

    Dataset mnist downloaded and prepared to /root/tensorflow_datasets/mnist/3.0.1. Subsequent calls will reuse this data.
    Downloading and preparing dataset fashion_mnist/3.0.1 (download: 29.45 MiB, generated: 36.42 MiB, total: 65.87 MiB) to /root/tensorflow_datasets/fashion_mnist/3.0.1...

    HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Dl Completed...', max=1.0, style=Progre…

    HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Dl Size...', max=1.0, style=ProgressSty…

    HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Extraction completed...', max=1.0, styl…

    HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))

    Shuffling and writing examples to /root/tensorflow_datasets/fashion_mnist/3.0.1.incomplete1LWAVV/fashion_mnist-train.tfrecord

    HBox(children=(FloatProgress(value=0.0, max=60000.0), HTML(value='')))

    HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))

    Shuffling and writing examples to /root/tensorflow_datasets/fashion_mnist/3.0.1.incomplete1LWAVV/fashion_mnist-test.tfrecord

    HBox(children=(FloatProgress(value=0.0, max=10000.0), HTML(value='')))

    Dataset fashion_mnist downloaded and prepared to /root/tensorflow_datasets/fashion_mnist/3.0.1. Subsequent calls will reuse this data.

[Chex](https://github.com/deepmind/chex) is a library of utilities for helping to write reliable JAX code.

Within `chex` you will find a `dataclass` object definition, which will automatically register new class instances into JAX, so you can easily apply JAX's tree utilities out of the box. We will use it to define a labelled data object type.

```
@chex.dataclass
class ContextualData():
  target: chex.Array
  context: chex.Array
```

```
# Here we provide some utilities for experiment and data visualization.

def gallery(array, ncols=None):
  """Rearrange an array of images into a tiled layout."""
  nindex, height, width, num_channels = array.shape
  if ncols is None:
    ncols = int(np.sqrt(nindex))
  nrows = int(np.ceil(nindex/ncols))
  pad = np.zeros((nrows*ncols-nindex, height, width, num_channels))
  array = np.concatenate([array, pad], axis=0)
  result = (array.reshape(nrows, ncols, height, width, num_channels)
            .swapaxes(1,2)
            .reshape(height*nrows, width*ncols, num_channels))
  return result

def imshow(x, title=''):
  """Shorthand for imshow."""
  plt.imshow(x[..., 0], cmap='gist_yarg', interpolation=None)
  plt.axis('off')
  plt.title(title)

def custom_after_subplot(ax: plt.Axes, group_name: str, x_label: str):
  """Disable Legend in LiveLossPlot (interactive Matplotlib)"""
  ax.set_title(group_name)
  ax.set_xlabel(x_label)
  ax.legend().set_visible(False)
```

# The datasets

In this tuturial we will train _conditional_ models.
As such, and as suggested by the `ContextualData` defintion, the dataset will provide targets and contexts. We will use two types of context, depending on the dataset.

- **Simple**: for FashionMNIST the context will be the object class label, in one-hot format.
- **Hard**: for MNIST digits we will make things a bit more interesting: the user will specify a simple function, mapping tuples of integers in the range `[0,9]` to an integer in `[0, 9]`, for example:
```
lambda i, j: (i + j) % 10
```
The context will be a tuple of images from MNIST whose labels will match the context integers, and the target will be an MNIST image of the corresponging
output label. Hance, the result will be a dataset whose conditional context is potentially very rich and challenging to capture.

```
# --------------------------------- #
# SIMPLE context dataset generation #
# --------------------------------- #

def get_fashion_mnist(batch_size, data_split='train', seed=1, conditional=True):
  def _preprocess(sample):
    image = tf.cast(sample["image"], tf.float32) / 255.0
    context = tf.one_hot(sample["label"], 10)
    # We can optionally make context constant, effectively making the dataset
    # unconditional. This can be useful for debuggin purposes.
    if not conditional:
      context *= 0
    return ContextualData(target=image, context=context)

  ds = fashion_mnist[data_split]
  ds = ds.map(map_func=_preprocess,
              num_parallel_calls=tf.data.experimental.AUTOTUNE)
  ds = ds.cache()
  ds = ds.shuffle(100000, seed=seed).repeat().batch(batch_size)
  return iter(tfds.as_numpy(ds))

# Visualize a sample of the data
ds = next(get_fashion_mnist(batch_size=12))
d = gallery(ds.target)
imshow(d, title='Fashion MNIST targets')
```

```
# ------------------------------- #
# HARD context dataset generation #
# ------------------------------- #

def get_raw_data(data_split='train'):
  def _preprocess(sample):
    image = tf.cast(sample["image"], tf.float32) / 255.0
    id = sample["label"]
    return image, id

  ds = mnist[data_split]
  ds = ds.map(map_func=_preprocess,
              num_parallel_calls=tf.data.experimental.AUTOTUNE)
  ds = ds.cache()
  ds = ds.shuffle(100000, seed=0).batch(2048)
  images, labels = next(iter(tfds.as_numpy(ds)))

  data_by_label = []
  for i in range(10):
    data_by_label.append(images[labels==i])
  min_num = min([d.shape[0] for d in data_by_label])
  return np.stack([d[:min_num] for d in data_by_label])

def get_mnist_digits(batch_size,
                     fn=None,
                     data_split='train',
                     max_num_examplars=None):
  """Instantiates a data generation function implementing the input fn in mnist.

  Args:
    batch_size: (int) batch size of the returned data
    fn: (python function) default sum of two digits modulo 10, function from n
        digits to a digit. The generator will return a batch of n + 1 digit
        images whose labels correspont to inputs and output of the function.
    data_split: (string) default 'train', crop of mnist to use for the data.
    max_num_examplars: (int) maximum number of exemplars to be used - even when
        not specified all digits will be represented by the same number of
        exemplars.

  Returns:
    Function of JAX PRNG returning samples from mnist related to each other via
    the user input fn function.
  """
  raw_data = get_raw_data(data_split)
  sz = raw_data.shape[1]
  if max_num_examplars is not None:
    sz = min(max_num_examplars, sz)

  if fn is None:
    # Examples
    # fn = lambda i, j: (i + j) % 10
    # fn = lambda i: (i + 1) % 10
    fn = lambda i: (i + 1) % 10

  num_inputs = len(inspect.signature(fn).parameters)
  def data_fn(prng):
    digit_prng, sample_prng = jax.random.split(prng, 2)
    digit_indices = jax.random.randint(
        digit_prng, shape=[batch_size, num_inputs], minval=0, maxval=10)
    digit_indices = digit_indices.split(num_inputs, axis=1)
    digit_indices += [fn(*digit_indices)]

    sample_indices = jax.random.randint(
        sample_prng, shape=[batch_size, num_inputs+1], minval=0, maxval=sz)
    sample_indices = sample_indices.split(num_inputs + 1, axis=1)
    samples = [raw_data[i[..., 0], j[..., 0], ...]
               for i, j in zip(digit_indices, sample_indices)]
    return samples

  def generator(key=None):
    if key is None:
      key = jax.random.PRNGKey(0)
    while True:
      key, sample_key = jax.random.split(key)
      data = data_fn(sample_key)
      yield ContextualData(target=data[-1], context=data[:-1])

  return generator

def tile_context(context, filler_value=.2):
  entries = [gallery(d) for d in data.context]
  entry_shape = entries[0].shape
  separator = np.full((entry_shape[0], 1, 1), filler_value)
  entries = [
    np.concatenate([e, separator], axis=1) for e in entries[:-1]] + entries[-1:]
  return np.concatenate(entries, axis=1)

# Visualize samples of target with hard context function
context_fn = lambda x, y, z: (x + y + z) % 10
data = next(get_mnist_digits(fn=context_fn, batch_size=16)())
sz = 5
plt.figure(figsize=(4 * sz, sz))
plt.subplot(122)
imshow(gallery(data.target),
       'Target with ' + get_code_as_string(context_fn))
plt.subplot(121)
imshow(tile_context(data.context) , 'Context')
```

For convenience, we define a dataset constructor function taking a `hard` flag switching between the MNIST and Fashion MNIST datasets, as well a dummy dataset
constructor that we will use for the purpose of retrieving shape information at graph construction time.

```
def get_dataset(batch_size, num_dev, hard=False, data_split='train'):
  # Instantiates the dataset adjusting the batch size to support training on
  # multiple devices.
  if hard:
    dataset = get_mnist_digits(
        batch_size=batch_size*num_dev,
        data_split=data_split)()
  else:
    dataset = get_fashion_mnist(
        batch_size=batch_size*num_dev,
        data_split=data_split)
  return dataset

def get_dummy_data(hard=False):
  # Returns an instance of the data to gather shape info.
  return next(get_dataset(batch_size=1, num_dev=1, hard=hard))
```

# Amortized variational inference (VAEs)

Consider a joint distribution $p(x, z)$ over a set of latent variables $z \in \mathcal{Z}$ and observed variable $x \in \mathcal{X}$ (for instance, the images of our dataset).

Making inference over the observed variable $x$ involves computing the posterior distribution $p(z|x) = p(x,z)/p(x)$ which is often intractable to compute, as the _marginal likelihood_ $p(x) = \int_z p(x, z)dz$ requires integrating over a potentially exponential number of configurations of $z$.

**Variational Inference (VI)** can be used to approximate the posterior $p(z|x)$ in a tractable fashion. VI casts the problem of computing the posterior as an optimization problem introducing a family of tractable (simpler) distributions $\mathcal{Q}$ parametrized by $\lambda$. The objective is to find the best approximation of the true posterior $q_{\lambda^*} \in \mathcal{Q}$ that minimizes the Kullback-Leibler (KL) divergence with the exact prosterior:
$$
q_{\lambda^*}(z) = \underset{q_{\lambda}}{arg min} \ \ D_{KL}(q_{\lambda}(z) || p(z|x))
$$

$q_{\lambda^*}(z)$ can serve as a proxy for the true posterior distribution. Note that the solution depends on the speciﬁc value of the observed (evidence) variables $x_i$ we are conditioning on, so computing the posterior requires solving an optimization problem for each sample independently.

In this tutorial, we use a much more efficient approach. Rather than solving an optimization process per data point, we can **amortize the cost of inference** by leveraging the power of function approximation and learn a deterministic mapping to predict the distributional variables as a function of $x$. Specifically, the posterior parameters for $x_i$ will be the output of a *learned* function $f_\theta(x_i)$, where $\theta$ are parameters shared across all data points.

For more information, see:
  * [Kingma and Welling, (2013), Auto-Encoding Variational Bayes](https://arxiv.org/abs/1312.6114)
  * [Kingma and Welling, (2019), An Introduction to Variational Autoencoders](https://arxiv.org/abs/1906.02691)

# Implement and train a (Conditional) Variational AutoEncoder

Variational AutoEncoders (VAEs) are a powerful class of deep generative models with latent variables, comprising of three main components: an encoder, a decoder and a prior.

The encoder, also called *inference* or *recognition model*, computes the approximate posterior distribution $q_{\phi}(z|x)$ of a given sample conditioned on $x$, while the decoder, or *generative model* reconstructs the sample starting from the latent variables $z$. In the conditional setting, the encoder and decoder will potentially also receive an additional conditioning input that we will refer to as _context_.

When we have a unconditional generative model, we generally don't have control over the specific outputs produced by the model, and the generated samples could be anything depending on the sampled latent variable. Perhaps, a more useful generative process should allow one to influence specific characteristics of the samples to generate, for instance, based on a context variable $c$, e.g. a label or category. For this reason, in this tutorial, we focus on Conditional VAEs (CVAEs) where both the encoder and the decoder are conditioned on a context variable $c$.

The main components of our model will be:

* $P^*$ is the true data distribution. We have some samples from this in the form of a dataset.
* $p(z)$ is a *prior* distribution over the latent space. In general the prior is simply $\mathcal{N}(0, 1)$ but in our case we will *learn* a conditional prior distribution $p(z|c)$ based on the context.
* $E(x, c)$ the encoder outputs distributions over the latent space $Z$, not just elements of it. The produced distribution is denoted $q_\phi(z|x, c)$ and is the (approximate) *posterior* distribution.
* $D(z, c)$ the decoder may be stochastic again, modeling the output distribution $p_\theta(x|z, c)$.

Now we go over the reconstruction and sampling process and finally motivate the losses used to train VAEs in this tutorial.

### Reconstruction
The process for reconstruction is:

1. Take $x, c \sim P^*$.
2. Encode it $E_\phi(x|c)$, yielding $q_\phi(z|x, c)$.
3. Sample a latent $z \sim q_\phi(z|x, c)$.
4. Decode the latent $p_\theta(\hat{x}|z, c) = D_\theta(z, c)$.
5. Sample a reconstruction: $\hat{x} \sim p_\theta(\hat{x}|z, c)$.

The prior has not showed up here, it plays a role in sampling.

Illustration of the conditional VAE architecture.
 Credits:  https://ijdykeman.github.io/ml/2016/12/21/cvae.html

### Sampling
The sampling process is:

1. Given a context $c$, sample a latent $z \sim p(z|c)$ from the conditional prior.
2. Decode the latent $p_\theta(x|z, c) = D_\theta(z, c)$.
3. Sample a reconstruction: $x \sim p_\theta(x|z, c)$.

In practice we usually use simple, parametrizable distributions in the encoder and decoder. More specifically:

**Encoder**
Each latent dimension is a (univariate) gaussian, parametrized by mean and standard deviation. Note, this is the same as a multivariate guassian over the latent space with a diagional covariance matrix.

**Decoder**
We will quantize the pixels to 0 and 1, which allows us to use a Bernoulli distribution per pixel to model it. Though for visualizations we will continue to use grayscale values.

## The Loss

We use maximum likelihood for training, that is, ideally we would like to maximize:

$$\mathbb{E}_{x,c \sim P^*}\log p_{\theta}(x|c).$$

Note that $p_{\theta}(x|c)$ is the marginal probability distribution $p_{\theta}(x|c) = \int p_\theta(x, z|c) dz$. We can rewrite this in familiar terms as $\int p_\theta(x|z,c) p(z|c) dz$. However, computing (and maximizing) the above marginal is computationally infeasible.

Instead, we can show

$$\log p_{\theta}(x|c) \ge \mathbb{E}_{z \sim q(z|x,c)} \big[\log p_\theta(x | z,c)\big] - \mathbb{KL}\big(q_\phi(z | x,c) || p(z|c)\big).$$

This right hand side is called the evidence lower bound (ELBO). Broadly speaking the term variational methods, like variational inference, refers to this technique of using an approximate posterior distribution and the ELBO; this is where Variational Autoencoder gets its name from too.

In order to try to maximize the likelihood, we maximize the ELBO instead. Recall from the lecture that under some conditions (that are not going to apply to us) the inequality is actually an equality. This yield the following loss used with Variational AutoEncoders:

$$ \mathcal{L}(x|c) = - \Big( \mathbb{E}_{z \sim q(z|x, c)} \big[\log p_\theta(x | z, c)\big] - \mathbb{KL}\big(q_\phi(z | x, c) || p(z|c)\big) \Big).$$

Observe that:
* The first term encourages the model to reconstruct the input faithfully. This part is similar to the Vanilla AutoEncoder.
* The second term can be seen as a *regularization term* of the encoder towards the prior.
* Encoder, Decoder and Prior are conditioned on the context $c$. Removing the conditioning corresponds to recovering the original VAE formulation.

(The formula contains an expectation; in practice that would be approximated with one or more samples.)

## Implementation
We start by defining the main VAE object class.

We want to decouple its interface from how its components are defined; for example at this point we don't really care if we are using neural networks to implement the mapping from inputs to posterior.

Some of the implementation will have to be `haiku` specific, but we will make an effort to restrict these details into the initialization functions.

```
def _transform(component):
  return hk.without_apply_rng(hk.transform(component))

# Sum-reduce on all dimensions but batch size
_batch_sum = jax.vmap(jnp.sum)

class VAE():
  """A (conditional) Variational Autoencoder.

  The class expects at construction time haiku modules implementing the
  conditional {encoder, decoder, prior}, as well as context projector module.

  The encoder, decoder and prior use the output of the context projector
  together with their other inputs to define distributions (posterior, output
  and conditional prior respectively), implemented as tfp.distributions.

  The context projector is used to map a potentially nested context to a single
  condition tensor.
  """

  def __init__(self, *, encoder, decoder, prior, context_projector):
    self._encoder = _transform(encoder)
    self._decoder = _transform(decoder)
    self._prior = _transform(prior)
    self._context_projector = _transform(context_projector)

  def init_params(self, prng, data):
    prng_encoder, prng_decoder, prng_prior, prng_proj = jax.random.split(
        prng, 4)

    # Initialize the context mapping.
    # ADD CODE BELOW
    # -----------------------
    context_projector_params = ...
    projected_context = ...
    # -----------------------

    # Initialize the conditional prior.
    # ADD CODE BELOW
    # -----------------------
    prior_params = ...
    z = ...
    # -----------------------

    # Initialize the encoder.
    # ADD CODE BELOW
    # -----------------------
    encoder_params = ...
    # -----------------------

    # Initialize the decoder.
    # ADD CODE BELOW
    # -----------------------
    decoder_params =  ...
    # -----------------------

    # Merge all the parameters into a single data structure.
    params = hk.data_structures.merge(
        context_projector_params, prior_params, encoder_params, decoder_params)
    return params

  def sample(self, params, prng, context, mean=True):
    prior_prng, decoder_prng = jax.random.split(prng)
    # Map the context.
    # ADD CODE BELOW
    # -----------------------
    projected_context = ...
    # -----------------------

    # Get the conditional prior distribution, and take a sample from it.
    # ADD CODE BELOW
    # -----------------------
    conditional_prior = ...
    z = ...
    # -----------------------

    # Get the conditional output distribution.
    # ADD CODE BELOW
    # -----------------------
    output_distribution = ...
    # -----------------------
    if mean:
      return output_distribution.mean()
    else:
      return output_distribution.sample(seed=decoder_prng)

  def sample_prior(self, params, prng, context):
    projected_context = self._context_projector.apply(params, context)
    conditional_prior = self._prior.apply(params, projected_context)
    return conditional_prior.sample(seed=prng)

  def decode(self, params, prng, z, context, mean=True):
    projected_context = self._context_projector.apply(params, context)
    output_distribution = self._decoder.apply(params, z, projected_context)
    if mean:
      return output_distribution.mean()
    else:
      return output_distribution.sample(seed=prng)

  def reconstruct(self, params, prng, x, context, mean=True):
    posterior_prng, decoder_prng = jax.random.split(prng)
    # Map the context.
    # ADD CODE BELOW
    # -----------------------
    projected_context = ...
    # -----------------------

    # Get the conditional posterior distribution p(z|x, context), and sample
    # from it.
    # ADD CODE BELOW
    # -----------------------
    posterior = ...
    z = ...
    # -----------------------

    # Get the conditional output distribution.
    # ADD CODE BELOW
    # -----------------------
    output_distribution = ...
    # -----------------------
    if mean:
      return output_distribution.mean()
    else:
      return output_distribution.sample(seed=decoder_prng)

  def stochastic_elbo(self, params, prng, x, context):
    """ELBO = log_p(x) - KL(q(z|x) || p(z))."""
    projected_context = self._context_projector.apply(params, context)
    z, kl = self._z_and_kl(params, prng, x, projected_context)
    output_distribution = self._decoder.apply(params, z, projected_context)
    # Sum reduce over signal domain (but not over batch!)
    log_p_x = _batch_sum(output_distribution.log_prob(x))
    elbo = log_p_x - kl
    # Assemble all the stats that we want to log in an extra output dictionary.
    extra_outputs = dict(kl=kl, log_p=log_p_x, elbo=elbo)
    return elbo, extra_outputs

  def kl(self, params, prng, x, context):
    projected_context = self._context_projector.apply(params, context)
    _, kl = self._z_and_kl(params, prng, x, projected_context)
    return kl

  def _z_and_kl(self, params, prng, x, projected_context):
    """Sample z and compute KL(q(z|x) || p(z)) = log_p(q(z|x)) - log_p(p(x))."""

    # Get the conditional prior, given the pre-projected context
    # ADD CODE BELOW
    # -----------------------
    prior = ...
    # -----------------------

    # Get the conditional posterior distribution p(z|x, context), and sample
    # from it.
    # ADD CODE BELOW
    # -----------------------
    posterior = ...
    z = ...
    # -----------------------

    # Compute the posterior log probability for the sampled z.
    # Note _batch_sum: what is this doing?
    log_q_z = _batch_sum(posterior.log_prob(z))

    # Compute prior log probability for the sampled z.
    # ADD CODE BELOW
    # -----------------------
    log_p_z = ...
    # -----------------------

    # Compute the KL (see formula)
    # ADD CODE BELOW
    # -----------------------
    kl = ...
    # -----------------------
    return z, kl

```

Here we define all the utilities we will need to implement Haiku modules for the prior, encoder and decoder required to instantiate a VAE.

```
def _make_positive(x):
  """Transforms elementwise unconstrained inputs into positive outputs."""
  # The offset is such that the output will be equal to 1 whenever the input is
  # equal to 0.
  offset = jnp.log(jnp.exp(1.) - 1.)
  return jax.nn.softplus(x + offset)

class DiagonalNormal(tfp.distributions.Normal):
  """Normal distribution with diagonal covariance."""

  def __init__(self, params, name='diagonal_normal'):
    if params.shape[-1] != 2:
      raise ValueError(
          f'The last dimension of `params` must be 2, got {params.shape[-1]}.')
    super().__init__(
        loc=params[..., 0], scale=_make_positive(params[..., 1]), name=name)

class ConditionalPrior(hk.Module):
  """A prior distribution whose parameters are computed by a conditioner."""

  def __init__(self, map_ctor, distribution, name='prior_net'):
    super().__init__(name=name)
    # This function will map: (bs,) + context shape --> (bs, num_latents, 2)
    self._map = map_ctor()
    self._distribution = distribution

  def __call__(self, context):
    return self._distribution(self._map(context))

class ConditionalEncoder(hk.Module):
  """A posterior distribution whose parameters are computed by a conditioner."""

  def __init__(self, map_ctor, distribution, name='posterior_net'):
    super().__init__(name=name)
    # This function will map inputs to the to posterior distribution parameters.
    self._map = map_ctor()
    self._distribution = distribution

  def __call__(self, x, context):
    # We assume that x is an image and the context is flat.
    chex.assert_rank(x, 4)
    chex.assert_rank(context, 2)

    # Tile the context in order to be concatenated to the input tensor x.
    # ADD CODE BELOW
    # -----------------------
    bs, height, width, _ = x.shape
    context = ...
    x_and_context = ...
    # -----------------------

    # Compute the posterior q(z|x, context)
    return self._distribution(self._map(x_and_context))

class ConditionalDecoder(hk.Module):
  """An output distribution whose parameters are computed by a conditioner."""

  def __init__(self, map_ctor, distribution, name='output_net'):
    super().__init__(name=name)
    # This function will map (z, context) to the to output distribution
    # parameters.
    self._map = map_ctor()
    self._distribution = distribution

  def __call__(self, z, context):
    chex.assert_equal_shape_prefix([z, context], -1)
    # Concatenate the context to the latents
    # ADD CODE BELOW
    # -------------------------------------------------------------
    z_and_context = ...
    # -------------------------------------------------------------
    return self._distribution(self._map(z_and_context))

```

Here we define how to assemble the VAE.

```
# -- Flat model
def get_model(num_latents=50, num_hiddens=500, data_shape=[28, 28, 1]):
  """Creates a fully-connected VAE model."""

  def _build_map_to_latent_dist_params_net():
    return hk.Sequential([
        hk.Flatten(),
        hk.nets.MLP([num_hiddens, num_hiddens, num_latents * 2]),
        hk.Reshape([num_latents, 2]),
    ])

  def encoder_fn(*args, **kwargs):
    return ConditionalEncoder(
        map_ctor=_build_map_to_latent_dist_params_net,
        distribution=DiagonalNormal)(*args, **kwargs)

  def prior_fn(*args, **kwargs):
    return ConditionalPrior(
        map_ctor=_build_map_to_latent_dist_params_net,
        distribution=DiagonalNormal)(*args, **kwargs)

  def _build_map_to_output_dist_params_net():
    return hk.Sequential([
        hk.Flatten(),
        hk.nets.MLP([num_hiddens, num_hiddens, np.prod(data_shape)]),
        hk.Reshape(data_shape),
    ])

  def decoder_fn(*args, **kwargs):
    return ConditionalDecoder(
        map_ctor=_build_map_to_output_dist_params_net,
        distribution=tfp.distributions.Bernoulli)(*args, **kwargs)

  def _concatenate(tensors):
    proj = hk.nets.MLP([num_hiddens, num_latents])
    data = []
    for x in jax.tree_leaves(tensors):
      if x.ndim > 2:
        x = proj(hk.Flatten()(x))
      # Note: this is not commutative!
      # What could we do to make this projection commutative?
      data.append(x)
    data = jnp.concatenate(data, axis=-1)
    return data

  def proj_fn(*args):
    """Maps context inputs to conditioning tensor."""
    return hk.to_module(_concatenate)('context_projector')(*args)

  return VAE(encoder=encoder_fn,
             decoder=decoder_fn,
             prior=prior_fn,
             context_projector=proj_fn)

```

# Training

## Shared hyper-parameters

```
# @title Shared hyper-parameters

# Model hyper-parameters: you can play around with these and see what changes!
BATCH_SIZE = 128        #@param {type:'integer'}
VIS_BATCH_SIZE = 64     #@param {type:'integer'}
NUM_LATENTS = 30        #@param {type:'integer'}
NUM_HIDDENS = 50        #@param {type:'integer'}
TRAINING_STEPS = 1000   #@param {type:'integer'}
USE_HARD_DATA = False   #@param {type:'boolean'}
NUM_DEV = len(jax.local_devices())

# Plot configs: no need to change them
PLOT_REFRESH_EVERY = 50   #@param {type:'integer'}
PLOT_EVERY = 5            #@param {type:'integer'}
```

```
model = get_model(num_latents=NUM_LATENTS, num_hiddens=NUM_HIDDENS)

# Once we have trained a model, we can explore its latent space by interpolating
# (and decoding) between latents.
def get_latent_interpolations(model, params, prng_key, num_steps, context):
  start_key, stop_key, decoder_key = jax.random.split(prng_key, 3)
  z_start = model.sample_prior(params, start_key, data.context)
  z_stop = model.sample_prior(params, stop_key, data.context)
  # To ensure that the interpolation is still likely under the Gaussian prior,
  # we use Gaussian interpolation - rather than linear interpolation.
  a = jnp.linspace(.0, 1.0, num_steps)
  a = jnp.expand_dims(a, axis=1)
  interpolations =  (jnp.sqrt(a) * z_start + jnp.sqrt(1 - a) * z_stop)
  ncols = int(jnp.sqrt(num_steps))
  context = jax.tree_map(
      lambda x: jnp.tile(x, [num_steps] +  [1] *  (x.ndim-1)), data.context)
  samples_from_interpolations = model.decode(
      params, decoder_key, interpolations, context, mean=True)
  return samples_from_interpolations
```

```
# These are utilities for shape manipulation to facilitate training using pmap.
@jax.jit
def format_data(data):
  return jax.tree_map(
      lambda x: x.reshape((NUM_DEV, BATCH_SIZE) + x.shape[1:]), data)

def setup_for_distributed_training(params, prng_key, opt_state, num_dev):
  broadcast = lambda x: jnp.broadcast_to(x, (num_dev,) + x.shape)
  (params, opt_state) = jax.tree_map(broadcast, (params, opt_state))
  prng_key = jax.random.split(prng_key, num_dev)
  return params, prng_key, opt_state

get_slice = lambda x: jax.tree_map(lambda x: x[0], x)

def reconstruct_and_sample(params, prng_key, model, data):
  prng_key = prng_key[0]
  params = get_slice(params)
  sample_key, reconstruction_key, prng = jax.random.split(prng_key, 3)
  reconstruction = model.reconstruct(
      params, reconstruction_key, data.target, data.context)
  sample = model.sample(params, sample_key, data.context)
  return data.target, reconstruction, sample
```

# Training a VAE optimizing ELBO

```
def get_elbo_loss_fn(model):
  def elbo_loss_fn(params, prng, data):
    elbo, stats = jax.tree_map(
        jnp.mean, model.stochastic_elbo(params, prng, data.target, data.context))
    return -elbo, stats
  return elbo_loss_fn

```

```
elbo_loss_fn = get_elbo_loss_fn(model)

# Initialize the model
prng_key = jax.random.PRNGKey(0)
params = model.init_params(prng_key, get_dummy_data(hard=USE_HARD_DATA))

# Instantiate and initialize the optimizer
elbo_optimizer = tx.adam(1e-2)
opt_state = elbo_optimizer.init(params)

# Setup the optimizer state and params for distributed training
params, prng_key, opt_state = setup_for_distributed_training(
    params, prng_key, opt_state, num_dev=NUM_DEV)

dataset = get_dataset(
    batch_size=BATCH_SIZE, num_dev=NUM_DEV, hard=USE_HARD_DATA)

# Define and pmap the update function
def elbo_update(params, prng_key, opt_state, data):
  loss_key, prng_key = jax.random.split(prng_key)

  # Use jax.value_and_grad to compute amd return all the outputs and the
  # gradients for elbo_loss_fn (tip: what does the has_aux kwargs do?)
  # ADD CODE BELOW
  # -------------------------------------------------------------
  loss_outputs, grads = ...
  # -------------------------------------------------------------

  # Reduce-mean the gradients across devices.
  # ADD CODE BELOW
  # -------------------------------------------------------------
  grads = ...
  # -------------------------------------------------------------

  # Perform an udpate step on the elbo_optimizer to produce parameter udpates,
  # them apply the updates of the model params.
  # ADD CODE BELOW
  # -------------------------------------------------------------
  raw_updates, opt_state = ...
  params = ...
  # -------------------------------------------------------------
  return params, prng_key, opt_state, loss_outputs
elbo_update = jax.pmap(elbo_update, axis_name='i', devices=jax.devices())

# -- Set up logging and interactive plotting
losses = []
kls = []
logps = []

plot = PlotLosses(
  groups={
    'log p(x)': ['log_p'],
    'KL': ['kl'],
    'negative ELBO': ['negative ELBO']},
  outputs=[MatplotlibPlot(max_cols=3, after_subplot=custom_after_subplot)],
  step_names='Iterations')
```

```
for step in range(TRAINING_STEPS):
  params, prng_key, opt_state, stats = elbo_update(
      params, prng_key, opt_state, format_data(next(dataset)))

  elbo = stats[1]['elbo'].mean()
  kl = stats[1]['kl'].mean()
  log_p = stats[1]['log_p'].mean()

  if (step + 1) % PLOT_EVERY == 0:
    plot.update({
        'negative ELBO': -elbo,
        'kl': kl,
        'log_p': log_p
    }, current_step=step)
  if (step + 1) % PLOT_REFRESH_EVERY == 0:
    plot.send()

  losses.append(elbo)
  kls.append(kl)
  logps.append(log_p)
```

## Visualize reconstructions and samples

```
targets, reconstructions, samples = reconstruct_and_sample(
    params, prng_key, model,
    next(get_dataset(batch_size=VIS_BATCH_SIZE, num_dev=1, hard=USE_HARD_DATA,
                     data_split='test')))

sz = 6
_ = plt.figure(figsize=((3*sz, 1*sz)))
plt.subplot(131)
imshow(gallery(targets), 'Targets')
plt.subplot(132)
imshow(gallery(reconstructions), 'Reconstructions')
plt.subplot(133)
imshow(gallery(samples), 'Conditional samples');
```

## Explore the latent space

```
data = get_dummy_data(hard=USE_HARD_DATA)
num_steps = 7 * 7
samples_from_interpolations = get_latent_interpolations(
    model, get_slice(params), jax.random.PRNGKey(1), num_steps, data.context)

sz = 6
_ = plt.figure(figsize=((2*sz, sz)))
plt.subplot(121)
imshow(data.target[0,...], 'Label example')
plt.subplot(122)
imshow(gallery(samples_from_interpolations), 'Latent space interpolation')
```

# Training a VAE with KL annealing

Stochastic Variational Inference can get stuck in local optima since there is a certain tension between the likelihood and the KL terms. To ease the optimization, we could use a schedule fpr $\beta$ where the KL coefficient is slowly annealed from $0$ to $1$ throughout training. This corresponds to weight more the reconstruction term at the beginning of the training and then move towards the fully stochastic variational objective. The modified objective becomes:

$$ \mathcal{L}(x|c) = - \Big( \mathbb{E}_{z \sim q(z|x, c)} \big[\log p_\theta(x | z, c)\big] - \beta \ \mathbb{KL}\big(q_\phi(z | x, c) || p(z|c)\big) \Big).$$

where the hyper-parameter $\beta$ is annealed throughout optimization.

**Tasks:**
- Modify the training objective with a linear KL annealing schedule.
- Look at the loss terms, what differences do you notice in their behaviours compared to the standard VAE objective?
- Look at the reconstructions and samples during training. What do you see?

For more information, see:
* [Bowman et al, (2015), Generating Sentences from a Continuous Space, (Section 3.1)](https://arxiv.org/abs/1511.06349)
* [Sønderby et al, (2016), Ladder variational Autoencoder, (Section 2.3)](https://arxiv.org/abs/1602.02282)
* [Burgess et al, (2018), Understanding disentangling in β-VAE](https://arxiv.org/abs/1804.03599)

```
def get_beta_vae_loss_fn(model):
  def beta_vae_loss_fn(params, prng, beta, data):
    _, loss_components = jax.tree_map(
        jnp.mean, model.stochastic_elbo(
            params, prng, data.target, data.context))
    # The beta-VAE training loss function is a function of the beta
    # hyper-parameter. Implement the loss using the components returned by
    # the stochastic_elbo of the model class.
    # elbo_loss_fn
    # ADD CODE BELOW
    # -------------------------------------------------------------
    loss = ...
    # -------------------------------------------------------------
    return loss, loss_components
  return beta_vae_loss_fn
```

```
prng_key = jax.random.PRNGKey(0)
params = model.init_params(prng_key, get_dummy_data(hard=USE_HARD_DATA))

beta_vae_loss_fn = get_beta_vae_loss_fn(model)

# Set up the optimzer and model like you did for the ELBO optimization
# ADD CODE BELOW
# -------------------------------------------------------------
beta_vae_optimizer = ...
opt_state = ...
params, prng_key, opt_state = ...
dataset = ...

def beta_vae_update(params, prng_key, opt_state, beta, data):
  ...
  return params, prng_key, opt_state, beta_vea_loss

beta_vae_update = ...
# -------------------------------------------------------------

# During the optimization of the beta-VAE we will be scheduling the value of
# beta. Use optax's polynomial schedule to set up a linear schedule, e.g.
# from 0.01 to 1.0 over TRAINING_STEPS//2 steps.
# ADD CODE BELOW
# -------------------------------------------------------------
beta_vae_current_step = ...
initial_beta, final_beta = ...
beta_schedule = ...
# -------------------------------------------------------------

beta_vae_losses = []
beta_vae_kls = []
beta_vae_logps = []

beta_vae_groups = {
  'log p(x)': ['log_p'],
  'kl': ['kl'],
  'negative ELBO': ['negative ELBO'],
  'beta': ['beta']
}

beta_vae_plot = PlotLosses(
    groups=beta_vae_groups,
    outputs=[MatplotlibPlot(max_cols=4, after_subplot=custom_after_subplot)],
    step_names='Iterations')
```

```
broadcast = lambda x: jnp.broadcast_to(x, (NUM_DEV,) + x.shape)

for _ in range(TRAINING_STEPS):
  beta_vae_current_step += 1
  beta = beta_schedule(beta_vae_current_step)
  # Since beta is varying during training, it cannot be statically compiled into
  # udpate function and needs to be broadcast prior to being passed to the
  # beta_vae_update.
  params, prng_key, opt_state, stats = beta_vae_update(
      params, prng_key, opt_state, broadcast(beta), format_data(next(dataset)))

  elbo = stats[1]['elbo'].mean()
  kl = stats[1]['kl'].mean()
  log_p = stats[1]['log_p'].mean()

  if (beta_vae_current_step + 1) % PLOT_EVERY == 0:
    beta_vae_plot.update({
        'negative ELBO': -elbo,
        'kl': kl,
        'log_p': log_p,
        'beta': beta
    }, current_step=beta_vae_current_step)
  if beta_vae_current_step % PLOT_REFRESH_EVERY == 0:
    beta_vae_plot.send()

  beta_vae_losses.append(elbo)
  beta_vae_kls.append(kl)
  beta_vae_logps.append(log_p)
```

## Visualize reconstructions and samples

```
targets, reconstructions, samples = reconstruct_and_sample(
    params, prng_key, model,
    next(get_dataset(batch_size=VIS_BATCH_SIZE, num_dev=1, hard=USE_HARD_DATA,
                     data_split='test')))

sz = 6
_ = plt.figure(figsize=((3*sz, 1*sz)))
plt.subplot(131)
imshow(gallery(targets), 'Targets')
plt.subplot(132)
imshow(gallery(reconstructions), 'Reconstructions')
plt.subplot(133)
imshow(gallery(samples), 'Conditional samples');
```

## Explore the latent space

```
data = get_dummy_data(hard=USE_HARD_DATA)
num_steps = 7 * 7
samples_from_interpolations = get_latent_interpolations(
    model, get_slice(params), jax.random.PRNGKey(1), num_steps, data.context)

sz = 6
_ = plt.figure(figsize=((2*sz, sz)))
plt.subplot(121)
imshow(data.target[0,...], 'Label example')
plt.subplot(122)
imshow(gallery(samples_from_interpolations), 'Latent space interpolation')
```

# Constrained optimization

Constrained optimization can be used to dynamically tune the relative weight of the likelihood and KL terms during optimization.
This removes the need to manually tune $\beta$, or create an optimization schedule, which can be problem specific.

The objective now becomes:

\begin{equation}
  \text{minimize } \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z)) \text{ such that }  \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x) \log p_\theta(x|z)} > \kappa
\end{equation}

This can be solved using Lagrange multipliers. The objective then becomes:

\begin{equation}
  \text{minimize }  \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z)) + \lambda  (\mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\kappa - \log p_\theta(x|z)))
\end{equation}

The difference compared to the KL annealing is that:

   * $\lambda$ is a learned parameter - it will be learned using stochastic gradient descent, like the network parameters.  The difference is that the lagrangian has to solve a maximization problem. You can see this intuitively: the graadient with respect to $\lambda$ in the objective above is $\mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z))$. If $ \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z))> 0$, the constraint is not being satisfied, so the value of the lagrangian needs to increase. This will be done by doing gradient ascent, instead of gradient descent. Note that for $\lambda$ to be a valid lagranian in a minimization problem, it has to be positive.
   * The practicioner has to specify the hyperparameter $\kappa$, which determines the reoncstruction quality of the model.
   * the coefficient is in front of the likelihood term, not the KL term. This is mainly for convenience, as it is easier to specify the hyperparameter $\kappa$ for the likelihood (reconstruction loss).

For more assumptions made by this method, see the Karush–Kuhn–Tucker conditions.

For more information, see:
  * http://bayesiandeeplearning.org/2018/papers/33.pdf

## Lagrange multipliers for constrained optimization

This is a reimplementation in JAX of the [constrained optimization tools](https://github.com/deepmind/sonnet/blob/master/sonnet/python/modules/optimization_constraints.py) found in Sonnet v1.

Note the peculiar implementation of the `clip` function, which is used to limit the valid range of the multipliers; this is used for example to force their values to be non-negative, or to limit their maximum values to control the interaction of the loss components during training.

In this implementation the gradients flowing through the clipping function are not set to 0 when above and below the thresholds, if the gradients would make the inputs move towards the valid range. This is particularly useful when setting a maximum valid value, since it allows for the multipliers to become smaller once the constraints are (eventually) satisfied.

```
@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3))
def clip(x, min_val, max_val, maximize=True):
  del maximize
  return jax.tree_map(lambda e: jnp.clip(e, min_val, max_val), x)

def clip_fwd(x, min_val, max_val, maximize):
  return clip(x, min_val, max_val, maximize), x

def clip_bwd(min_val, max_val, maximize, x, co_tangents):
  zeros = jax.tree_map(jnp.zeros_like, co_tangents)
  if min_val is not None:
    get_mask = lambda x, v, t: (x < v) & (t < 0 if maximize else t > 0)
    mask = jax.tree_multimap(get_mask, x, min_val, co_tangents)
    co_tangents = jax.tree_multimap(jnp.where, mask, zeros, co_tangents)
  if max_val is not None:
    get_mask = lambda x, v, t: (x > v) & (t > 0 if maximize else t < 0)
    mask = jax.tree_multimap(get_mask, x, max_val, co_tangents)
    co_tangents = jax.tree_multimap(jnp.where, mask, zeros, co_tangents)
  return co_tangents,

clip.defvjp(clip_fwd, clip_bwd)

class LagrangeMultiplier(hk.Module):
  """Lagrange Multiplier module for constrained optimization."""

  def __init__(self,
               shape=(),
               initializer=1.0,
               maximize=True,
               valid_range=None,
               name='lagrange_multiplier'):
    super().__init__(name=name)
    self._shape = shape
    if callable(initializer):
      self._initializer = initializer
    else:
      self._initializer = lambda *args: jnp.ones(*args) * initializer
    self._maximize = maximize
    self._valid_range = valid_range if valid_range is not None else (0.0, None)
    assert self._valid_range[0] >= 0

  def __call__(self):
    lag_mul = hk.get_parameter('w', self._shape, init=self._initializer)
    lag_mul = clip(lag_mul, *self._valid_range, maximize=self._maximize)
    return lag_mul
```

## Training a VAE using GECO

```
def get_geco_loss(prng, model, kappa, valid_range, init_lambda):

  def constraint_term(x, target):
    lag_mul = LagrangeMultiplier(
        shape=x.shape, valid_range=valid_range, initializer=init_lambda)()
    return jnp.sum(lag_mul * (x - target)), lag_mul
  constraint_term = _transform(constraint_term)
  lagmul_params = constraint_term.init(prng, jnp.ones(()), jnp.ones(()))

  def geco_loss(params, prng, data):
    """Loss using constrained optimization as in GECO."""
    # Compute the KL term (make sure to average across the batch!)
    # ADD CODE BELOW
    # -----------------------
    kl = ...
    # -----------------------

    # Reconstruct the inputs and compute the reconstruction error (e.g. MSE)
    # ADD CODE BELOW
    # -----------------------
    reconstruction = ...
    reconstruction_err = ...
    # -----------------------

    # If you didn't use MSE before, you need to adjust the kappa term in the
    # constraint definition.
    constraint_satisfaction, lag_mul = constraint_term.apply(
        params, reconstruction_err, np.prod(data.target.shape[1:]) * kappa**2)
    loss = constraint_satisfaction + kl
    metrics = {
        'loss': loss,
        'mse': jnp.mean((reconstruction - data.target)**2),
        'kl': kl,
        'lag_mul': lag_mul}
    return loss, metrics

  return geco_loss, lagmul_params
```

```
prng_key = jax.random.PRNGKey(0)
kappa = 0.18
valid_range = (0, 4.0)
init_lambda = 2.0

params = model.init_params(prng_key, get_dummy_data(hard=USE_HARD_DATA))
geco_loss_fn, lagmul_params = get_geco_loss(
    prng_key, model, kappa, valid_range, init_lambda=init_lambda)
```

## Using multiple optimizers jointly

We will now see how we can optimize joinlty the model parameters `params` and the Lagrange multipliers `lagmul_params`.

We want to update the Lagrange multipliers via stochastic gradient *ascent*; for example we could consider using the logic in `tx.sgd` with a *negative* learning rate. The model params will be instead optimized using ADAM, just like in the previous cells.

In the next cell you will find a handy wrapper to capture multiple optimizers into a single optax object.

What is left to do is:

1.   Merge all the parameters into a single tree; you can do so using the `hk.data_structures.merge` utility in Haiku:

  ```
  params = hk.data_structures.merge(params, lagmul_params)
  ```
2.   Specify how to filter the resulting `params` structure to select which variables should be mapped to which optimizer. Tip: you can filter the Lagrange multipliers with:
  ```
  multiplier_filter = lambda m, n, p: 'lagrange' in m
  ```
  and the model variables with:
  ```
  model_params_filter = lambda m, n, p: 'lagrange' not in m
```

```
def multi_opt(**filter_optimizer_map):
  """Wraps multiple optimizers within an object with the optax interface.

  Args:
    **filter_optimizer_map: kwargs used to map optimizer names to
    (predicate, optimizer) tuples, where predicate is a function which will be
    passed to haiku.data_structure.filter in order to select which variables
    are to be updated by the corresponding optimizer.

  Returns:
    An optax.InitUpdate tuple.
  """
  def filter_(predicate, params):
    if params is None:
      return None
    return hk.data_structures.filter(predicate, params)

  merge_ = hk.data_structures.merge

  def _init(params):
    opt_state = dict()
    for opt_name, (predicate, opt) in filter_optimizer_map.items():
      opt_state[opt_name] = opt.init(filter_(predicate, params))
    return opt_state

  def _update(updates, state, params=None):
    new_updates = {}
    new_state = dict()
    for opt_name, (predicate, opt) in filter_optimizer_map.items():
      update, new_state[opt_name] = opt.update(
          filter_(predicate, updates), state[opt_name],
          filter_(predicate, params))
      new_updates = merge_(new_updates, update)
    return new_updates, new_state

  return tx.InitUpdate(_init, _update)
```

```
# Merge them using hk.data_structures.merge
# function.
# ADD CODE BELOW
# -----------------------
params = ...
# -----------------------

# Optimize the Lagrange multipliers via stochastic gradient ascent.
# Optimize the model params will be optimized using ADAM.
# ADD CODE BELOW
# -----------------------
geco_optimizer = ...
opt_state = ...
# -----------------------

params, prng_key, opt_state = setup_for_distributed_training(
    params, prng_key, opt_state, num_dev=NUM_DEV)

dataset = get_dataset(
    batch_size=BATCH_SIZE, num_dev=NUM_DEV, hard=USE_HARD_DATA)

def geco_update(params, prng_key, opt_state, data):
  loss_key, prng_key = jax.random.split(prng_key)
  geco_loss, grads = jax.value_and_grad(
      geco_loss_fn, has_aux=True)(params, loss_key, data)
  grads = jax.lax.pmean(grads, axis_name='i')
  raw_updates, opt_state = geco_optimizer.update(grads, opt_state)
  params = tx.apply_updates(params, raw_updates)
  return params, prng_key, opt_state, geco_loss
geco_update = jax.pmap(geco_update, axis_name='i', devices=jax.devices())

geco_current_step = 0
mses = []
kls = []
lag_muls = []

geco_groups = {
  'mse': ['mse'],
  'kl': ['kl'],
  'lag_mul': ['lag_mul'],
}

geco_plot = PlotLosses(
    groups=geco_groups,
    outputs=[MatplotlibPlot(max_cols=4, after_subplot=custom_after_subplot)],
    step_names='Iterations')

```

```
for _ in range(TRAINING_STEPS):
  geco_current_step += 1
  params, prng_key, opt_state, stats = geco_update(
      params, prng_key, opt_state, format_data(next(dataset)))
  mse = stats[1]['mse'].mean()
  kl = stats[1]['kl'].mean()
  lag_mul = stats[1]['lag_mul'].mean()

  geco_plot.update({
      'mse': mse,
      'kl': kl,
      'lag_mul': lag_mul
  })
  if geco_current_step % PLOT_REFRESH_EVERY == 0:
    geco_plot.send()

  mses.append(mse)
  kls.append(kl)
  lag_muls.append(lag_mul)
```

## Visualize reconstructions and samples

```
targets, reconstructions, samples = reconstruct_and_sample(
    params, prng_key, model,
    next(get_dataset(batch_size=VIS_BATCH_SIZE, num_dev=1, hard=USE_HARD_DATA,
                     data_split='test')))

sz = 6
_ = plt.figure(figsize=((3*sz, 1*sz)))
plt.subplot(131)
imshow(gallery(targets), 'Targets')
plt.subplot(132)
imshow(gallery(reconstructions), 'Reconstructions')
plt.subplot(133)
imshow(gallery(samples), 'Conditional samples');
```

## Explore the latent space

```
dummy_data = get_dummy_data(hard=USE_HARD_DATA)
num_steps = 7 * 7
samples_from_interpolations = get_latent_interpolations(
    model, get_slice(params), jax.random.PRNGKey(0), num_steps, dummy_data.context)

sz = 6
_ = plt.figure(figsize=((2*sz, sz)))
plt.subplot(121)
imshow(dummy_data.target[0,...], 'Label example')
plt.subplot(122)
imshow(gallery(samples_from_interpolations), 'Latent space interpolation')
```


# Homework 4: Word-level entailment with neural networks

```python
__author__ = "Christopher Potts"
__version__ = "CS224u, Stanford, Spring 2020"
```

## Contents

1. [Overview](#Overview)
1. [Set-up](#Set-up)
1. [Data](#Data)
  1. [Edge disjoint](#Edge-disjoint)
  1. [Word disjoint](#Word-disjoint)
1. [Baseline](#Baseline)
  1. [Representing words: vector_func](#Representing-words:-vector_func)
  1. [Combining words into inputs: vector_combo_func](#Combining-words-into-inputs:-vector_combo_func)
  1. [Classifier model](#Classifier-model)
  1. [Baseline results](#Baseline-results)
1. [Homework questions](#Homework-questions)
  1. [Hypothesis-only baseline [2 points]](#Hypothesis-only-baseline-[2-points])
  1. [Alternatives to concatenation [2 points]](#Alternatives-to-concatenation-[2-points])
  1. [A deeper network [2 points]](#A-deeper-network-[2-points])
  1. [Your original system [3 points]](#Your-original-system-[3-points])
1. [Bake-off [1 point]](#Bake-off-[1-point])

## Overview

The general problem is word-level natural language inference.

Training examples are pairs of words $(w_{L}, w_{R}), y$ with $y = 1$ if $w_{L}$ entails $w_{R}$, otherwise $0$.

The homework questions below ask you to define baseline models for this and develop your own system for entry in the bake-off, which will take place on a held-out test-set distributed at the start of the bake-off. (Thus, all the data you have available for development is available for training your final system before the bake-off begins.)

## Set-up

See [the first notebook in this unit](nli_01_task_and_data.ipynb) for set-up instructions.

```python
from collections import defaultdict
import json
import numpy as np
import os
import pandas as pd
from torch_shallow_neural_classifier import TorchShallowNeuralClassifier
import nli
import utils
```

```python
DATA_HOME = 'data'

NLIDATA_HOME = os.path.join(DATA_HOME, 'nlidata')

wordentail_filename = os.path.join(
    NLIDATA_HOME, 'nli_wordentail_bakeoff_data.json')

GLOVE_HOME = os.path.join(DATA_HOME, 'glove.6B')
```

## Data

I've processed the data into two different train/test splits, in an effort to put some pressure on our models to actually learn these semantic relations, as opposed to exploiting regularities in the sample.

* `edge_disjoint`: The `train` and `dev` __edge__ sets are disjoint, but many __words__ appear in both `train` and `dev`.
* `word_disjoint`: The `train` and `dev` __vocabularies are disjoint__, and thus the edges are disjoint as well.

These are very different problems. For `word_disjoint`, there is real pressure on the model to learn abstract relationships, as opposed to memorizing properties of individual words.

```python
with open(wordentail_filename) as f:
    wordentail_data = json.load(f)
```

The outer keys are the  splits plus a list giving the vocabulary for the entire dataset:

```python
wordentail_data.keys()
```

### Edge disjoint

```python
wordentail_data['edge_disjoint'].keys()
```

This is what the split looks like; all three have this same format:

```python
wordentail_data['edge_disjoint']['dev'][: 5]
```

Let's test to make sure no edges are shared between `train` and `dev`:

```python
nli.get_edge_overlap_size(wordentail_data, 'edge_disjoint')
```

As we expect, a *lot* of vocabulary items are shared between `train` and `dev`:

```python
nli.get_vocab_overlap_size(wordentail_data, 'edge_disjoint')
```

This is a large percentage of the entire vocab:

```python
len(wordentail_data['vocab'])
```

Here's the distribution of labels in the `train` set. It's highly imbalanced, which will pose a challenge for learning. (I'll go ahead and reveal that the `dev` set is similarly distributed.)

```python
def label_distribution(split):
    return pd.DataFrame(wordentail_data[split]['train'])[1].value_counts()
```

```python
label_distribution('edge_disjoint')
```

### Word disjoint

```python
wordentail_data['word_disjoint'].keys()
```

In the `word_disjoint` split, no __words__ are shared between `train` and `dev`:

```python
nli.get_vocab_overlap_size(wordentail_data, 'word_disjoint')
```

Because no words are shared between `train` and `dev`, no edges are either:

```python
nli.get_edge_overlap_size(wordentail_data, 'word_disjoint')
```

The label distribution is similar to that of `edge_disjoint`, though the overall number of examples is a bit smaller:

```python
label_distribution('word_disjoint')
```

## Baseline

Even in deep learning, __feature representation is vital and requires care!__ For our task, feature representation has two parts: representing the individual words and combining those representations into a single network input.

### Representing words: vector_func

Let's consider two baseline word representations methods:

1. Random vectors (as returned by `utils.randvec`).
1. 50-dimensional GloVe representations.

```python
def randvec(w, n=50, lower=-1.0, upper=1.0):
    """Returns a random vector of length `n`. `w` is ignored."""
    return utils.randvec(n=n, lower=lower, upper=upper)
```

```python
# Any of the files in glove.6B will work here:

glove_dim = 50

glove_src = os.path.join(GLOVE_HOME, 'glove.6B.{}d.txt'.format(glove_dim))

# Creates a dict mapping strings (words) to GloVe vectors:
GLOVE = utils.glove2dict(glove_src)

def glove_vec(w):
    """Return `w`'s GloVe representation if available, else return
    a random vector."""
    return GLOVE.get(w, randvec(w, n=glove_dim))
```

### Combining words into inputs: vector_combo_func

Here we decide how to combine the two word vectors into a single representation. In more detail, where `u` is a vector representation of the left word and `v` is a vector representation of the right word, we need a function `vector_combo_func` such that `vector_combo_func(u, v)` returns a new input vector `z` of dimension `m`. A simple example is concatenation:

```python
def vec_concatenate(u, v):
    """Concatenate np.array instances `u` and `v` into a new np.array"""
    return np.concatenate((u, v))
```

`vector_combo_func` could instead be vector average, vector difference, etc. (even combinations of those) – there's lots of space for experimentation here; [homework question 2](#Alternatives-to-concatenation-[1-point]) below pushes you to do some exploration.

### Classifier model

For a baseline model, I chose `TorchShallowNeuralClassifier`:

```python
net = TorchShallowNeuralClassifier(hidden_dim=50, max_iter=100)
```

### Baseline results

The following puts the above pieces together, using `vector_func=glove_vec`, since `vector_func=randvec` seems so hopelessly misguided for `word_disjoint`!

```python
word_disjoint_experiment = nli.wordentail_experiment(
    train_data=wordentail_data['word_disjoint']['train'],
    assess_data=wordentail_data['word_disjoint']['dev'],
    model=net,
    vector_func=glove_vec,
    vector_combo_func=vec_concatenate)
```

## Homework questions

Please embed your homework responses in this notebook, and do not delete any cells from the notebook. (You are free to add as many cells as you like as part of your responses.)

### Hypothesis-only baseline [2 points]

During our discussion of SNLI and MultiNLI, we noted that a number of research teams have shown that hypothesis-only baselines for NLI tasks can be remarkably robust. This question asks you to explore briefly how this baseline effects the 'edge_disjoint' and 'word_disjoint' versions of our task.

For this problem, submit code the following:

1. A `vector_combo_func` function called `hypothesis_only` that simply throws away the premise, using the unmodified hypothesis (second) vector as its representation of the example.

1. Code for looping over the two conditions 'word_disjoint' and 'edge_disjoint' and the two `vector_combo_func` values `vec_concatenate` and `hypothesis_only`, calling `nli.wordentail_experiment` to train on the conditions 'train' portion and assess on its 'dev' portion, with `glove_vec` as the `vector_func`. So that the results are consistent, use an `sklearn.linear_model.LogisticRegression` with default parameters as the model.

1. Print out the percentage-wise increase in macro-F1 over the `hypothesis_only` baseline that `vec_concatenate` delivers for each of the two conditions. For example, if `hypothesis_only` returns 0.52 for condition `C` and  `vec_concatenate` delivers 0.75 for `C`, then you'd report a ((0.75 /  0.52) - 1) * 100  = 44.23 percent increase for `C`. The values you need are stored in the dictionary returned by `nli.wordentail_experiment`, with key 'macro-F1'. Please round the percentages to two digits.

```python
##### YOUR CODE HERE

```

### Alternatives to concatenation [2 points]

We've so far just used vector concatenation to represent the premise and hypothesis words. This question asks you to explore a simple alternative.

For this problem, submit code the following:

1. A new potential value for `vector_combo_func` that does something different from concatenation. Options include, but are not limited to, element-wise addition, difference, and multiplication. These can be combined with concatenation if you like.
1. Include a use of `nli.wordentail_experiment` in the same configuration as the one in [Baseline results](#Baseline-results) above, but with your new value of `vector_combo_func`.

```python
##### YOUR CODE HERE

```

### A deeper network [2 points]

It is very easy to subclass `TorchShallowNeuralClassifier` if all you want to do is change the network graph: all you have to do is write a new `define_graph`. If your graph has new arguments that the user might want to set, then you should also redefine `__init__` so that these values are accepted and set as attributes.

For this question, please subclass `TorchShallowNeuralClassifier` so that it defines the following graph:

$$\begin{align}
h_{1} &= xW_{1} + b_{1} \\
r_{1} &= \textbf{Bernoulli}(1 - \textbf{dropout_prob}, n) \\
d_{1} &= r_1 * h_{1} \\
h_{2} &= f(d_{1}) \\
h_{3} &= h_{2}W_{2} + b_{2}
\end{align}$$

Here, $r_{1}$ and $d_{1}$ define a dropout layer: $r_{1}$ is a random binary vector of dimension $n$, where the probability of a value being $1$ is given by $1 - \textbf{dropout_prob}$. $r_{1}$ is multiplied element-wise by our first hidden representation, thereby zeroing out some of the values. The result is fed to the user's activation function $f$, and the result of that is fed through another linear layer to produce $h_{3}$. (Inside `TorchShallowNeuralClassifier`, $h_{3}$ is the basis for a softmax classifier, so no activation function is applied to it.)

For comparison, using this notation, `TorchShallowNeuralClassifier` defines the following graph:

$$\begin{align}
h_{1} &= xW_{1} + b_{1} \\
h_{2} &= f(h_{1}) \\
h_{3} &= h_{2}W_{2} + b_{2}
\end{align}$$

The following code starts this sub-class for you, so that you can concentrate on `define_graph`. Be sure to make use of `self.dropout_prob`

For this problem, submit just your completed  `TorchDeepNeuralClassifier`. You needn't evaluate it, though we assume you will be keen to do that!

```python
import torch.nn as nn

class TorchDeepNeuralClassifier(TorchShallowNeuralClassifier):
    def __init__(self, dropout_prob=0.7, **kwargs):
        self.dropout_prob = dropout_prob
        super().__init__(**kwargs)

    def define_graph(self):
        """Complete this method!

        Returns
        -------
        an `nn.Module` instance, which can be a free-standing class you
        write yourself, as in `torch_rnn_classifier`, or the outpiut of
        `nn.Sequential`, as in `torch_shallow_neural_classifier`.

        """
        ##### YOUR CODE HERE

```

### Your original system [3 points]

This is a simple dataset, but our focus on the 'word_disjoint' condition ensures that it's a challenging one, and there are lots of modeling strategies one might adopt.

You are free to do whatever you like. We require only that your system differ in some way from those defined in the preceding questions. They don't have to be completely different, though. For example, you might want to stick with the model but represent examples differently, or the reverse.

Keep in mind that, for the bake-off evaluation, the 'edge_disjoint' portions of the data are off limits. You can, though, train on the combination of the 'word_disjoint' 'train' and 'dev' portions. You are free to use different pretrained word vectors and the like. Please do not introduce additional entailment datasets into your training data, though.

Please embed your code in this notebook so that we can rerun it.

In the cell below, please provide a brief technical description of your original system, so that the teaching team can gain an understanding of what it does. This will help us to understand your code and analyze all the submissions to identify patterns and strategies.

```python
# Enter your system description in this cell.
# Please do not remove this comment.

```

## Bake-off [1 point]

The goal of the bake-off is to achieve the highest macro-average F1 score on __word_disjoint__, on a test set that we will make available at the start of the bake-off. The announcement will go out on the discussion forum. To enter, you'll be asked to run `nli.bake_off_evaluation` on the output of your chosen `nli.wordentail_experiment` run.

The cells below this one constitute your bake-off entry.

The rules described in the [Your original system](#Your-original-system-[3-points]) homework question are also in effect for the bake-off.

Systems that enter will receive the additional homework point, and systems that achieve the top score will receive an additional 0.5 points. We will test the top-performing systems ourselves, and only systems for which we can reproduce the reported results will win the extra 0.5 points.

Late entries will be accepted, but they cannot earn the extra 0.5 points. Similarly, you cannot win the bake-off unless your homework is submitted on time.

The announcement will include the details on where to submit your entry.

```python
# Enter your bake-off assessment code into this cell.
# Please do not remove this comment.
##### YOUR CODE HERE

```

```python
# On an otherwise blank line in this cell, please enter
# your macro-avg f1 value as reported by the code above.
# Please enter only a number between 0 and 1 inclusive.
# Please do not remove this comment.

```