```python
%matplotlib notebook
%matplotlib inline
import math
import matplotlib.pyplot as plt
```

# Nuclear Power Economics and Fuel Management

## Syllabus

Throughout the semester, you can always find the syllabus online at [https://github.com/katyhuff/npre412/blob/master/syllabus/syllabus.pdf](https://github.com/katyhuff/npre412/blob/master/syllabus/syllabus.pdf).

```python
from IPython.display import IFrame
IFrame("../syllabus/syllabus.pdf", width=1000, height=1000)
```

```python
from IPython.display import IFrame
IFrame("http://katyhuff.github.io", width=1000, height=700)
```

## Assessment

My goal (and, hopefully your goal) is for you to learn this material. If I have done my job right, your grade in this class will reflect just that -- how much you have learned. To do this for many learning types, your comprehension of the readings will be assessed with quizzes, your ability to apply what you've learned from class on your own will be assessed with the homeworks and projects, and your wholistic retention of the material will be assessed with tests.

### Monday Background Reading Assignments

Rather than introducing you to concepts during class, I think our time is better spent if we focus on exploring those concepts through demonstration and discussion together. **This 'active learning' educational strategy is backed by science, but is also just more respectful of your ability to learn things on your own.** Therefore, the lectures will assume you have studied the background materials ahead of time. This will include book sections, government reports, videos, and other resources. You will be expected to study the material outside of class before the start of each week. This will include book sections, government reports, videos, and other resources. I recommend you take notes on this material as it may be part of the tests.

**On Monday of each week I'll assign a list of material. You'll have 7 days to study that material before we start covering those concepts in class.**

### Monday Quizzes

To help me calibrate the in-class discussion, a weekly quiz will assess your comprehension of the background material. The quizzes can be taken online through [Compass2g](https://compass2g.illinois.edu) at any time during the week, but they must be completed by Monday at 10am, 7 days after the material was assigned.

### Friday Homework Assignments

Homeworks will be assigned each Friday concerning the material covered that week. You will have 7 days to do the homework, so it will be due at 10am on the following Friday. You'll notice that my office hours are also on Friday. This is intentional, because I feel office hours are most effective if you come to them after handing in your homework to discuss the parts that you didn't get.

### Projects

The class will involve a longer project only assigned to graduate students.

### Tests

The midterms will take place in class. They will be independent of one another. The final will be comprehensive.

### Participation

I will notice when you are not in class, but attendance won't directly affect your grades. It may, however, indirectly affect your grades. If you miss something I demonstrate in class, you'll have a lot more trouble proving that you've learned it.

## How to get an A

My dear friend, mathematician Kathryn Mann, has a great summary of [how to get an A in her class.](https://math.berkeley.edu/~kpmann/getanA.pdf). Everything she says about her math classes is true for this class as well. You should expect to spend 3 hours outside of class for every hour you spend in class. So, for a 3 credit class, you'll need to spend 3 hours a week in class and 9 hours outside of class on the coursework. If you find you're spending much less or much more time on this class, please let me know.

## Late Work

**Late work has a halflife of 1 hour.** That is, adjusted for lateness, your grade $G(t)$ is a decaying percentage of the raw grade $G_0$. An assignment turned in $t$ hours late will recieve a grade according to the following relation:

$$
\begin{align}
        G(t) &= G_0e^{-\lambda t}\\
\end{align}
$$

where

$$
\begin{align}
        G(t) &= \mbox{grade adjusted for lateness}\\
        G_0 &= \mbox{raw grade}\\
        \lambda &= \frac{ln(2)}{t_{1/2}} = \mbox{decay constant} \\
        t &= \mbox{time elapsed since due [hours]}\\
        t_{1/2} &= 1 = \mbox{half-life [hours]} \\
\end{align}
$$

```python
import math
def late_grade(hours_late, grade=100, half_life=1):
    """This function describes how much credit you will get for late work"""
    lam = math.log(2)/half_life
    return grade*math.exp(-lam*hours_late)
```

```python
# This code plots how much credit you'll get over time
import numpy as np
y = np.arange(24)
x = np.arange(24)
for h in range(0,24):
    x[h] = h
    y[h] = late_grade(h)

# creates a figure and axes with matplotlib
fig, ax = plt.subplots()
scatter = plt.scatter(x, y, color='blue', s=y*20, alpha=0.4)
ax.plot(x, y, color='red')

# adds labels to the plot
ax.set_ylabel('Percent of Grade Earned')
ax.set_xlabel('Hours Late')
ax.set_title('Grade Decay')

# adds tooltips
import mpld3
labels = ['{0}% earned'.format(i) for i in y]
tooltip = mpld3.plugins.PointLabelTooltip(scatter, labels=labels)
mpld3.plugins.connect(fig, tooltip)

mpld3.display()
```

```python
print("If you turn in your homework an hour late, you'll get ", round(late_grade(1),2), "% credit.")
print("If you turn in your homework six hours late, you'll get ", round(late_grade(6),2), "% credit.")
print("If you turn in your homework a day late, you'll get ", round(late_grade(24),2), "% credit.")
print("If you turn in your homework two days late, you'll get ", round(late_grade(48),2), "% credit.")
print("If you turn in your homework three days late, you'll get ", round(late_grade(72),2), "% credit.")
```

    If you turn in your homework an hour late, you'll get  50.0 % credit.
    If you turn in your homework six hours late, you'll get  1.56 % credit.
    If you turn in your homework a day late, you'll get  0.0 % credit.
    If you turn in your homework two days late, you'll get  0.0 % credit.
    If you turn in your homework three days late, you'll get  0.0 % credit.

**There will be no negotiation about late work except in the case of absence documented by an absence letter from the Dean of Students.** The university policy for requesting such a letter is [here](http://studentcode.illinois.edu/article1_part5_1-501.html) . Please note that such a letter is appropriate for many types of conflicts, but that religious conflicts require special early handling. In accordance with university policy, students seeking an excused absence for religious reasons should complete the Request for Accommodation for Religious Observances Form, which can be found on the Office of the Dean of Students website. The student should submit this form to the instructor and the Office of the Dean of Students by the end of the second week of the course to which it applies.

## Communications

Things to try when you have a question:

- Be Persistent: [Try just one more time.](https://s-media-cache-ak0.pinimg.com/736x/03/54/ce/0354ce58a7a4308edcc46dd9238e12d7.jpg)
- Google: [You might be surprised at its depth.](https://devhumor.com/content/uploads//images/April2016/google-errors.jpg)
- Piazza: Try this first, your student colleagues probably know the answer.
- TA email: A quick question can usually be answered by your TA via email.
- TA office hours: Your TA is there for you at a regularly scheduled time.
- Prof. email: Questions not appropriate for your TA can be directed to me.
- Prof. office hours: I will be in my office once a week for your convenience.
- Prof. appointment: For private matters or when office hours conflict with your schedule, schedule an appointment with me.

### A note on email

[Email tips for dealing with fussy professor types.](https://medium.com/@lportwoodstacer/how-to-email-your-professor-without-being-annoying-af-cf64ae0e4087)

## Python, IPython, Jupyter, git, and the Notebooks

Rather than reading equations off of slides, I will display lecture notes, equations, and images in "Jupyter notebooks" like this one. Sometimes, I will call them by their old name, "IPython notebooks," but I'm talking about "Jupyter notebooks". Interleaved with the course notes, we will often write small functions in the Python programming language to represent the equations we are talking about. This will allow you to interact with the math, changing variables, modifying the models, and exploring the parameter space.

### But I don't know Python!

*You don't have to know Python to take this class.* However, you will need to learn a little along the way. I will provide lots of example code to support your completion of homework assignments and I will never ask you to write functioning code as part of any written exam. Programming is really hard without the internet.

### Exercises

Watch for blocks titled **Exercise** in the notebooks. Those mark moments when I will ask you, during class, to try something out, explore an equation, or arrive at an answer. These are short and are not meant to be difficult. They exist to quickly solidify an idea before we move on to the next one. I will often randomly call on students (with a random number generator populated with the enrollment list) to give solutions to the exercises, so **a failure to show up and participate will be noticed.**

### Installing Python, IPython, Jupyter, git, and the Notebooks

Because engaging in the exercises will be really helpful for you to study, you'll should try to gain access to a computer equipped with Python (a version greater than 3.0) and a basic set of scientific python libraries. If you have a computer already, I encourage you to install [anaconda](https://www.continuum.io/downloads).

These notebooks are stored "in the cloud," which is to say that they are stored on someone else's computers. Those computers are servers at GitHub, a sometimes silly but also very important company in the beautiful city of San Francisco. GitHub stores "git repositories" which are collections of files that are "version controlled" by the program "git." This is a lot to keep track of, and I won't require that you learn git to participate in this class. However, I strongly recommend using git and GitHub and to keep track of your research code. So, I encourage you to use git to access the notebooks.

**More information about the things you might want to install can be found in the [README](https://github.com/katyhuff/npre412/blob/master/README.md).**


##### Copyright 2020 The TensorFlow Authors.

```
#@title 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
#
# https://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.
```

# Quantum Convolutional Neural Network

    View on TensorFlow.org

    Run in Google Colab

    View source on GitHub

    Download notebook

This tutorial implements a simplified Quantum Convolutional Neural Network (QCNN), a proposed quantum analogue to a classical convolutional neural network that is also *translationally invariant*.

This example demonstrates how to detect certain properties of a quantum data source, such as a quantum sensor or a complex simulation from a device. The quantum data source being a cluster state that may or may not have an excitation—what the QCNN will learn to detect (The dataset used in the paper was SPT phase classification).

## Setup

```
!pip install tensorflow==2.3.1
```

Install TensorFlow Quantum:

```
!pip install tensorflow-quantum
```

Now import TensorFlow and the module dependencies:

```
import tensorflow as tf
import tensorflow_quantum as tfq

import cirq
import sympy
import numpy as np

# visualization tools
%matplotlib inline
import matplotlib.pyplot as plt
from cirq.contrib.svg import SVGCircuit
```

## 1. Build a QCNN

### 1.1 Assemble circuits in a TensorFlow graph

TensorFlow Quantum (TFQ) provides layer classes designed for in-graph circuit construction. One example is the `tfq.layers.AddCircuit` layer that inherits from `tf.keras.Layer`. This layer can either prepend or append to the input batch of circuits, as shown in the following figure.

The following snippet uses this layer:

```
qubit = cirq.GridQubit(0, 0)

# Define some circuits.
circuit1 = cirq.Circuit(cirq.X(qubit))
circuit2 = cirq.Circuit(cirq.H(qubit))

# Convert to a tensor.
input_circuit_tensor = tfq.convert_to_tensor([circuit1, circuit2])

# Define a circuit that we want to append
y_circuit = cirq.Circuit(cirq.Y(qubit))

# Instantiate our layer
y_appender = tfq.layers.AddCircuit()

# Run our circuit tensor through the layer and save the output.
output_circuit_tensor = y_appender(input_circuit_tensor, append=y_circuit)
```

Examine the input tensor:

```
print(tfq.from_tensor(input_circuit_tensor))
```

And examine the output tensor:

```
print(tfq.from_tensor(output_circuit_tensor))
```

While it is possible to run the examples below without using `tfq.layers.AddCircuit`, it's a good opportunity to understand how complex functionality can be embedded into TensorFlow compute graphs.

### 1.2 Problem overview

You will prepare a *cluster state* and train a quantum classifier to detect if it is "excited" or not. The cluster state is highly entangled but not necessarily difficult for a classical computer. For clarity, this is a simpler dataset than the one used in the paper.

For this classification task you will implement a deep MERA-like QCNN architecture since:

1.   Like the QCNN, the cluster state on a ring is translationally invariant.
2.   The cluster state is highly entangled.

This architecture should be effective at reducing entanglement, obtaining the classification by reading out a single qubit.

An "excited" cluster state is defined as a cluster state that had a `cirq.rx` gate applied to any of its qubits. Qconv and QPool are discussed later in this tutorial.

### 1.3 Building blocks for TensorFlow

One way to solve this problem with TensorFlow Quantum is to implement the following:

1. The input to the model is a circuit tensor—either an empty circuit or an X gate on a particular qubit indicating an excitation.
2. The rest of the model's quantum components are constructed with `tfq.layers.AddCircuit` layers.
3. For inference a `tfq.layers.PQC` layer is used. This reads $\langle \hat{Z} \rangle$ and compares it to a label of 1 for an excited state, or -1 for a non-excited state.

### 1.4 Data
Before building your model, you can generate your data. In this case it's going to be excitations to the cluster state (The original paper uses a more complicated dataset). Excitations are represented with `cirq.rx` gates. A large enough rotation is deemed an excitation and is labeled `1` and a rotation that isn't large enough is labeled `-1` and deemed not an excitation.

```
def generate_data(qubits):
    """Generate training and testing data."""
    n_rounds = 20  # Produces n_rounds * n_qubits datapoints.
    excitations = []
    labels = []
    for n in range(n_rounds):
        for bit in qubits:
            rng = np.random.uniform(-np.pi, np.pi)
            excitations.append(cirq.Circuit(cirq.rx(rng)(bit)))
            labels.append(1 if (-np.pi / 2) <= rng <= (np.pi / 2) else -1)

    split_ind = int(len(excitations) * 0.7)
    train_excitations = excitations[:split_ind]
    test_excitations = excitations[split_ind:]

    train_labels = labels[:split_ind]
    test_labels = labels[split_ind:]

    return tfq.convert_to_tensor(train_excitations), np.array(train_labels), \
        tfq.convert_to_tensor(test_excitations), np.array(test_labels)
```

You can see that just like with regular machine learning you create a training and testing set to use to benchmark the model. You can quickly look at some datapoints with:

```
sample_points, sample_labels, _, __ = generate_data(cirq.GridQubit.rect(1, 4))
print('Input:', tfq.from_tensor(sample_points)[0], 'Output:', sample_labels[0])
print('Input:', tfq.from_tensor(sample_points)[1], 'Output:', sample_labels[1])
```

### 1.5 Define layers

Now define the layers shown in the figure above in TensorFlow.

#### 1.5.1 Cluster state

The first step is to define the cluster state using Cirq, a Google-provided framework for programming quantum circuits. Since this is a static part of the model, embed it using the `tfq.layers.AddCircuit` functionality.

```
def cluster_state_circuit(bits):
    """Return a cluster state on the qubits in `bits`."""
    circuit = cirq.Circuit()
    circuit.append(cirq.H.on_each(bits))
    for this_bit, next_bit in zip(bits, bits[1:] + [bits[0]]):
        circuit.append(cirq.CZ(this_bit, next_bit))
    return circuit
```

Display a cluster state circuit for a rectangle of cirq.GridQubits:

```
SVGCircuit(cluster_state_circuit(cirq.GridQubit.rect(1, 4)))
```

#### 1.5.2 QCNN layers

Define the layers that make up the model using the Cong and Lukin QCNN paper. There are a few prerequisites:

* The one- and two-qubit parameterized unitary matrices from the Tucci paper.
* A general parameterized two-qubit pooling operation.

```
def one_qubit_unitary(bit, symbols):
    """Make a Cirq circuit enacting a rotation of the bloch sphere about the X,
    Y and Z axis, that depends on the values in `symbols`.
    """
    return cirq.Circuit(
        cirq.X(bit)**symbols[0],
        cirq.Y(bit)**symbols[1],
        cirq.Z(bit)**symbols[2])

def two_qubit_unitary(bits, symbols):
    """Make a Cirq circuit that creates an arbitrary two qubit unitary."""
    circuit = cirq.Circuit()
    circuit += one_qubit_unitary(bits[0], symbols[0:3])
    circuit += one_qubit_unitary(bits[1], symbols[3:6])
    circuit += [cirq.ZZ(*bits)**symbols[6]]
    circuit += [cirq.YY(*bits)**symbols[7]]
    circuit += [cirq.XX(*bits)**symbols[8]]
    circuit += one_qubit_unitary(bits[0], symbols[9:12])
    circuit += one_qubit_unitary(bits[1], symbols[12:])
    return circuit

def two_qubit_pool(source_qubit, sink_qubit, symbols):
    """Make a Cirq circuit to do a parameterized 'pooling' operation, which
    attempts to reduce entanglement down from two qubits to just one."""
    pool_circuit = cirq.Circuit()
    sink_basis_selector = one_qubit_unitary(sink_qubit, symbols[0:3])
    source_basis_selector = one_qubit_unitary(source_qubit, symbols[3:6])
    pool_circuit.append(sink_basis_selector)
    pool_circuit.append(source_basis_selector)
    pool_circuit.append(cirq.CNOT(control=source_qubit, target=sink_qubit))
    pool_circuit.append(sink_basis_selector**-1)
    return pool_circuit
```

To see what you created, print out the one-qubit unitary circuit:

```
SVGCircuit(one_qubit_unitary(cirq.GridQubit(0, 0), sympy.symbols('x0:3')))
```

And the two-qubit unitary circuit:

```
SVGCircuit(two_qubit_unitary(cirq.GridQubit.rect(1, 2), sympy.symbols('x0:15')))
```

And the two-qubit pooling circuit:

```
SVGCircuit(two_qubit_pool(*cirq.GridQubit.rect(1, 2), sympy.symbols('x0:6')))
```

##### 1.5.2.1 Quantum convolution

As in the Cong and Lukin paper, define the 1D quantum convolution as the application of a two-qubit parameterized unitary to every pair of adjacent qubits with a stride of one.

```
def quantum_conv_circuit(bits, symbols):
    """Quantum Convolution Layer following the above diagram.
    Return a Cirq circuit with the cascade of `two_qubit_unitary` applied
    to all pairs of qubits in `bits` as in the diagram above.
    """
    circuit = cirq.Circuit()
    for first, second in zip(bits[0::2], bits[1::2]):
        circuit += two_qubit_unitary([first, second], symbols)
    for first, second in zip(bits[1::2], bits[2::2] + [bits[0]]):
        circuit += two_qubit_unitary([first, second], symbols)
    return circuit
```

Display the (very horizontal) circuit:

```
SVGCircuit(
    quantum_conv_circuit(cirq.GridQubit.rect(1, 8), sympy.symbols('x0:15')))
```

##### 1.5.2.2 Quantum pooling

A quantum pooling layer pools from $N$ qubits to $\frac{N}{2}$ qubits using the two-qubit pool defined above.

```
def quantum_pool_circuit(source_bits, sink_bits, symbols):
    """A layer that specifies a quantum pooling operation.
    A Quantum pool tries to learn to pool the relevant information from two
    qubits onto 1.
    """
    circuit = cirq.Circuit()
    for source, sink in zip(source_bits, sink_bits):
        circuit += two_qubit_pool(source, sink, symbols)
    return circuit
```

Examine a pooling component circuit:

```
test_bits = cirq.GridQubit.rect(1, 8)

SVGCircuit(
    quantum_pool_circuit(test_bits[:4], test_bits[4:], sympy.symbols('x0:6')))
```

### 1.6 Model definition

Now use the defined layers to construct a purely quantum CNN. Start with eight qubits, pool down to one, then measure $\langle \hat{Z} \rangle$.

```
def create_model_circuit(qubits):
    """Create sequence of alternating convolution and pooling operators
    which gradually shrink over time."""
    model_circuit = cirq.Circuit()
    symbols = sympy.symbols('qconv0:63')
    # Cirq uses sympy.Symbols to map learnable variables. TensorFlow Quantum
    # scans incoming circuits and replaces these with TensorFlow variables.
    model_circuit += quantum_conv_circuit(qubits, symbols[0:15])
    model_circuit += quantum_pool_circuit(qubits[:4], qubits[4:],
                                          symbols[15:21])
    model_circuit += quantum_conv_circuit(qubits[4:], symbols[21:36])
    model_circuit += quantum_pool_circuit(qubits[4:6], qubits[6:],
                                          symbols[36:42])
    model_circuit += quantum_conv_circuit(qubits[6:], symbols[42:57])
    model_circuit += quantum_pool_circuit([qubits[6]], [qubits[7]],
                                          symbols[57:63])
    return model_circuit

# Create our qubits and readout operators in Cirq.
cluster_state_bits = cirq.GridQubit.rect(1, 8)
readout_operators = cirq.Z(cluster_state_bits[-1])

# Build a sequential model enacting the logic in 1.3 of this notebook.
# Here you are making the static cluster state prep as a part of the AddCircuit and the
# "quantum datapoints" are coming in the form of excitation
excitation_input = tf.keras.Input(shape=(), dtype=tf.dtypes.string)
cluster_state = tfq.layers.AddCircuit()(
    excitation_input, prepend=cluster_state_circuit(cluster_state_bits))

quantum_model = tfq.layers.PQC(create_model_circuit(cluster_state_bits),
                               readout_operators)(cluster_state)

qcnn_model = tf.keras.Model(inputs=[excitation_input], outputs=[quantum_model])

# Show the keras plot of the model
tf.keras.utils.plot_model(qcnn_model,
                          show_shapes=True,
                          show_layer_names=False,
                          dpi=70)
```

### 1.7 Train the model

Train the model over the full batch to simplify this example.

```
# Generate some training data.
train_excitations, train_labels, test_excitations, test_labels = generate_data(
    cluster_state_bits)

# Custom accuracy metric.
@tf.function
def custom_accuracy(y_true, y_pred):
    y_true = tf.squeeze(y_true)
    y_pred = tf.map_fn(lambda x: 1.0 if x >= 0 else -1.0, y_pred)
    return tf.keras.backend.mean(tf.keras.backend.equal(y_true, y_pred))

qcnn_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
                   loss=tf.losses.mse,
                   metrics=[custom_accuracy])

history = qcnn_model.fit(x=train_excitations,
                         y=train_labels,
                         batch_size=16,
                         epochs=25,
                         verbose=1,
                         validation_data=(test_excitations, test_labels))
```

```
plt.plot(history.history['loss'][1:], label='Training')
plt.plot(history.history['val_loss'][1:], label='Validation')
plt.title('Training a Quantum CNN to Detect Excited Cluster States')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()
```

## 2. Hybrid models

You don't have to go from eight qubits to one qubit using quantum convolution—you could have done one or two rounds of quantum convolution and fed the results into a classical neural network. This section explores quantum-classical hybrid models.

### 2.1 Hybrid model with a single quantum filter

Apply one layer of quantum convolution, reading out $\langle \hat{Z}_n \rangle$ on all bits, followed by a densely-connected neural network.

#### 2.1.1 Model definition

```
# 1-local operators to read out
readouts = [cirq.Z(bit) for bit in cluster_state_bits[4:]]

def multi_readout_model_circuit(qubits):
    """Make a model circuit with less quantum pool and conv operations."""
    model_circuit = cirq.Circuit()
    symbols = sympy.symbols('qconv0:21')
    model_circuit += quantum_conv_circuit(qubits, symbols[0:15])
    model_circuit += quantum_pool_circuit(qubits[:4], qubits[4:],
                                          symbols[15:21])
    return model_circuit

# Build a model enacting the logic in 2.1 of this notebook.
excitation_input_dual = tf.keras.Input(shape=(), dtype=tf.dtypes.string)

cluster_state_dual = tfq.layers.AddCircuit()(
    excitation_input_dual, prepend=cluster_state_circuit(cluster_state_bits))

quantum_model_dual = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_dual)

d1_dual = tf.keras.layers.Dense(8)(quantum_model_dual)

d2_dual = tf.keras.layers.Dense(1)(d1_dual)

hybrid_model = tf.keras.Model(inputs=[excitation_input_dual], outputs=[d2_dual])

# Display the model architecture
tf.keras.utils.plot_model(hybrid_model,
                          show_shapes=True,
                          show_layer_names=False,
                          dpi=70)
```

#### 2.1.2 Train the model

```
hybrid_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
                     loss=tf.losses.mse,
                     metrics=[custom_accuracy])

hybrid_history = hybrid_model.fit(x=train_excitations,
                                  y=train_labels,
                                  batch_size=16,
                                  epochs=25,
                                  verbose=1,
                                  validation_data=(test_excitations,
                                                   test_labels))
```

```
plt.plot(history.history['val_custom_accuracy'], label='QCNN')
plt.plot(hybrid_history.history['val_custom_accuracy'], label='Hybrid CNN')
plt.title('Quantum vs Hybrid CNN performance')
plt.xlabel('Epochs')
plt.legend()
plt.ylabel('Validation Accuracy')
plt.show()
```

As you can see, with very modest classical assistance, the hybrid model will usually converge faster than the purely quantum version.

### 2.2 Hybrid convolution with multiple quantum filters

Now let's try an architecture that uses multiple quantum convolutions and a classical neural network to combine them.

#### 2.2.1 Model definition

```
excitation_input_multi = tf.keras.Input(shape=(), dtype=tf.dtypes.string)

cluster_state_multi = tfq.layers.AddCircuit()(
    excitation_input_multi, prepend=cluster_state_circuit(cluster_state_bits))

# apply 3 different filters and measure expectation values

quantum_model_multi1 = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_multi)

quantum_model_multi2 = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_multi)

quantum_model_multi3 = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_multi)

# concatenate outputs and feed into a small classical NN
concat_out = tf.keras.layers.concatenate(
    [quantum_model_multi1, quantum_model_multi2, quantum_model_multi3])

dense_1 = tf.keras.layers.Dense(8)(concat_out)

dense_2 = tf.keras.layers.Dense(1)(dense_1)

multi_qconv_model = tf.keras.Model(inputs=[excitation_input_multi],
                                   outputs=[dense_2])

# Display the model architecture
tf.keras.utils.plot_model(multi_qconv_model,
                          show_shapes=True,
                          show_layer_names=True,
                          dpi=70)
```

#### 2.2.2 Train the model

```
multi_qconv_model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
    loss=tf.losses.mse,
    metrics=[custom_accuracy])

multi_qconv_history = multi_qconv_model.fit(x=train_excitations,
                                            y=train_labels,
                                            batch_size=16,
                                            epochs=25,
                                            verbose=1,
                                            validation_data=(test_excitations,
                                                             test_labels))
```

```
plt.plot(history.history['val_custom_accuracy'][:25], label='QCNN')
plt.plot(hybrid_history.history['val_custom_accuracy'][:25], label='Hybrid CNN')
plt.plot(multi_qconv_history.history['val_custom_accuracy'][:25],
         label='Hybrid CNN \n Multiple Quantum Filters')
plt.title('Quantum vs Hybrid CNN performance')
plt.xlabel('Epochs')
plt.legend()
plt.ylabel('Validation Accuracy')
plt.show()
```


[](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
print("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)]
```

# 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 Reinforcemente Learning setting we might want to explicitely lable 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)

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

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

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

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

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

def unbatch(data):
  flattened_data, tree_def = jax.tree_flatten(data)
  # ADD CODE BELOW
  # -----------------------
  split_data = ...
  unbatched_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))}]
```

## 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 avarage 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")
```

[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')
```