We will need to import some helper code, so we need to run this

```python
import os
import sys
nb_dir = os.path.split(os.getcwd())[0]
if nb_dir not in sys.path:
    sys.path.append(nb_dir)
```

# Colab

We will need to download some data for this notebook, so if you are using [colab](https://colab.research.google.com), set the `using_colab` flag below to `True` in order to clone our [github repo](https://github.com/probabll/dgm4nlp).

```python
using_colab = False
!ls
```

```python
if using_colab:
  !rm -fr dgm4nlp sst
  !git clone https://github.com/vitutorial/exercises.git
  !cp -R exercises/SST/* ./
  !ls
```

Now we can start our lab.

```python
import torch
from torch import nn
import torch.nn.functional as F
# CPU should be fine for this lab
device = torch.device('cuda:0')
#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
from collections import OrderedDict
import numpy as np
```

# Sentiment Classification

We are going to augment a sentiment classifier with a layer of discrete latent variables which will help us improve the model's interpretability. But first, let's quickly review the baseline task.

In sentiment classification, we have some text input $x = \langle x_1, \ldots, x_n \rangle$, e.g. a sentence or short paragraph, which expresses a certain sentiment $y$, i.e. one of $K$ classes, towards a subject (e.g. a film or a product).

We can learn a sentiment classifier by learning a categorical distribution over classes for a given input:

\begin{align}
Y|x &\sim \text{Cat}(f(x; \theta))
\end{align}

where the Categorical pmf is $\text{Cat}(y|\pi) = \pi_y$.

A categorical distribution over $K$ classes is parameterised by a $K$-dimensional probability vector, here we use a neural network $f$ to map from the input to this probability vector. Technically we say *a neural network parameterises our model*, that is, it computes the parameters of our categorical observation model. The figure below is a graphical depiction of the model: circled nodes are random variables (a shaded node is an observed variable), uncircled nodes are deterministic, a plate indicates multiple draws.

The neural network (NN) $f(\cdot; \theta)$ has parameters of its own, i.e. the weights of the various architecture blocks used, which we denote generically by $\theta$.

Suppose we have a dataset $\mathcal D = \{(x^{(1)}, y^{(1)}), \ldots, (x^{(N)}, y^{(N)})\}$ containing $N$ i.i.d. observations. Then we can use the log-likelihood function
\begin{align}
\mathcal L(\theta|\mathcal D) &= \sum_{k=1}^{N} \log P(y^{(k)}|x^{(k)}, \theta) \\
&= \sum_{k=1}^{N} \log \text{Cat}(y^{(k)}|f(x^{(k)}; \theta))
\end{align}
 to estimate $\theta$ by maximisation:
 \begin{align}
 \theta^\star = \arg\max_{\theta \in \Theta} \mathcal L(\theta|\mathcal D) ~ .
 \end{align}

We can use stochastic gradient-ascent to find a local optimum of $\mathcal L(\theta|\mathcal D)$, which only requires a gradient estimate:

\begin{align}
\nabla_\theta \mathcal L(\theta|\mathcal D) &= \sum_{k=1}^{|\mathcal D|} \nabla_\theta  \log P(y^{(k)}|x^{(k)}, \theta) \\
&= \sum_{k=1}^{|\mathcal D|} \frac{1}{N} N \nabla_\theta  \log P(y^{(k)}|x^{(k)}, \theta)  \\
&= \mathbb E_{\mathcal U(1/N)} \left[ N \nabla_\theta  \log P(y^{(K)}|x^{(K)}, \theta) \right]  \\
&\overset{\text{MC}}{\approx} \frac{N}{M} \sum_{m=1}^M \nabla_\theta  \log P(y^{(k_m)}|x^{(k_m)}, \theta) \\
&\text{where }K_m \sim \mathcal U(1/N)
\end{align}

This is a Monte Carlo (MC) estimate of the gradient computed on $M$ data points selected uniformly at random from $\mathcal D$.

For as long as $f$ remains differentiable wrt to its inputs and parameters, we can rely on automatic differentiation to obtain gradient estimates.

In what follows we show how to design $f$ and how to extend this basic model to a latent-variable model.

## Data

We provide you some code to load the data (see `sst.sstutil.examplereader`). Play with the snippet below and inspect a few training instances:

```python
from sstutil import examplereader, Vocabulary, load_glove

# Let's load the data into memory.
print("Loading data")
train_data = list(examplereader('data/sst/train.txt'))
dev_data = list(examplereader('data/sst/dev.txt'))
test_data = list(examplereader('data/sst/test.txt'))

print("train", len(train_data))
print("dev", len(dev_data))
print("test", len(test_data))

print('\n# Examples')
example = dev_data[0]
print("First dev example:", example)
print("First dev example tokens:", example.tokens)
print("First dev example label:", example.label)
```

## Architecture

The function $f$ conditions on a high-dimensional input (i.e. text), so we need to convert it to continuous real vectors. This is the job of an *encoder*.

**Embedding Layer**

The first step is to convert the words in $x$ to vectors, which in this lab we will do with a pre-trained embedding layer (we will use GloVe).

We will denote the embedding of the $i$th word of the input by:

\begin{equation}
\mathbf x_i = \text{glove}(x_i)
\end{equation}

**Encoder Layer**

In this lab, an encoder takes a sequence of input vectors $\mathbf x_1^n$, each $I$-dimensional, and produces a sequence of output vectors $\mathbf t_1^n$, each $O$-dimensional and a summary vector $\mathbf h \in \mathbb R^O$:

\begin{equation}
    \mathbf t_1^n, \mathbf h = \text{encoder}(\mathbf x_1^n; \theta_{\text{enc}})
\end{equation}

where we use $\theta_{\text{enc}}$ to denote the subset of parameters in $\theta$ that are specific to this encoder block.

*Remark:* in practice for a correct batched implementation, our encoders also take a mask matrix and a vector of lengths.

Examples of encoding functions can be a feed-forward NN (with an aggregator based on sum or average/max pooling) or a recurrent NN (e.g. an LSTM/GRU). Other architectures are also possible.

**Output Layer**

From our summary vector $\mathbf h$, we need to parameterise a categorical distribution over $K$ classes, thus we use

\begin{align}
f(x; \theta) &= \text{softmax}(\text{dense}_K(\mathbf h; \theta_{\text{output}}))
\end{align}

where $\text{dense}_K$ is a dense layer with $K=5$ outputs and $\theta_{\text{output}}$ corresponds to its parameters (weight matrix and bias vector). Note that we need to use the softmax activation function in order to guarantee that the output of $f$ is a normalised probability vector.

## Implementation

To leave an indication of the shape of tensors in the code, we use the following convention

```python
[B, T, D]
```

where `B` stands for `batch_size`, `T` stands for `time` (or rather *maximum sequence length*), and `D` is the size of the representation.

Consider the following abstract Encoder class:

```python
class Encoder(nn.Module):
    """
    An Encoder for us is a function that
      1. transforms a sequence of I-dimensional vectors into a sequence of O-dimensional vectors
      2. summarises a sequence of I-dimensional vectors into one O-dimensional vector

    """

    def __init__(self):
        super(Encoder, self).__init__()

    def forward(self, inputs, mask, lengths):
        """
        The input is a batch-first tensor of token ids. Here is an example:

        Example of inputs (though rather than words, we have word ids):
            INPUTS                     MASK       LENGTHS
            [the nice cat -PAD-]    -> [1 1 1 0]  [3]
            [the nice dog running]  -> [1 1 1 1]  [4]

        Note that:
              mask =  inputs == 1
              lengths = mask.sum(dim=-1)

        :param inputs: [B, T, I]
        :param mask: [B, T]
        :param lengths: [B]
        :returns: [B, T, O], [B, O]
            where the first tensor is the transformed input
            and the second tensor is a summary of all inputs
        """
        pass

```

Let's start easy, implement a *bag of words* encoder:

```python
class BagOfWordsEncoder(Encoder):
    """
    This encoder does not transform the input sequence,
     and its summary output is just a sum.
    """

    pass
```

```python
# SOLUTION
class BagOfWordsEncoder(Encoder):
    """
    This encoder does not transform the input sequence,
     and its summary output is just a sum.
    """

    def __init__(self):
        super(BagOfWordsEncoder, self).__init__()

    def forward(self, inputs, mask, lengths, **kwargs):
        return inputs, (inputs * mask.unsqueeze(-1).float()).sum(dim=1)
```

```python
def test_bow_encoder(encoder: BagOfWordsEncoder):
    """
    Use this to assert a few things, such as output shapes.
    """
    import numpy as np
    batch_size = 2
    max_seq_len = 10
    dim = 12
    inputs = torch.rand([batch_size, max_seq_len, dim])
    mask = torch.ones([batch_size, max_seq_len]).long()
    mask[-1,-1] = 0
    lengths = mask.sum(-1).long()
    outputs, h = encoder(inputs, mask, lengths)
    # General tests: output shape
    assert outputs.size(0) == batch_size, "Wrong value for size(0). Do you have a batch dimension?"
    assert outputs.size(1) == max_seq_len, "Wrong value for size(1). Do you have a time dimension?"
    assert outputs.size(2) == dim, "Wrong value for size(2). Do you have an embedding dimension?"
    np.testing.assert_equal(h[0].numpy(), outputs.sum(1)[0].numpy(), "Incorrect composition (did you sum the inputs?)")
    assert (h[1].numpy() != outputs.sum(1)[1].numpy()).sum() > 0, "Incorrect composition (have applied the mask?)"
test_bow_encoder(BagOfWordsEncoder())
```

You can also consider implementing

* a feed-forward encoder with average pooling
* and a biLSTM encoder

but these are certainly optional.

```python
class FFEncoder(Encoder):
    """
    A typical feed-forward NN with tanh hidden activations.
    """

    def __init__(self, input_size, output_size,
                 activation=None,
                 hidden_sizes=[],
                 aggregator='avg',
                 dropout=0.5):
        """
        :param input_size: int
        :param output_size: int
        :param hidden_sizes: list of integers (dimensionality of hidden layers)
        :param aggregator: 'sum' or 'avg'
        :param dropout: dropout rate
        """
        pass

    def forward(self, x, mask, lengths):
        """
        :param x: sequence of word embeddings, shape [B, T, I]
        :param mask: byte mask that is 0 for invalid positions, shape [B, T]
        :param lengths: the lengths of each input sequence [B]
        :return:
            outputs [B, T, O]
            sum/avg pooling [B, O]
        """
        pass
```

```python
#SOLUTION

class FFEncoder(Encoder):
    """
    A typical feed-forward NN with tanh hidden activations.
    """

    def __init__(self, input_size, output_size,
                 activation=None,
                 hidden_sizes=[],
                 aggregator='sum',
                 dropout=0.5):
        """
        :param input_size: int
        :param output_size: int
        :param hidden_sizes: list of integers (dimensionality of hidden layers)
        :param aggregator: 'sum' or 'avg'
        :param dropout: dropout rate
        """
        super(FFEncoder, self).__init__()
        layers = []
        if hidden_sizes:
            for i, size in enumerate(hidden_sizes):
                if dropout > 0.:
                  layers.append(('dropout%d' % i, nn.Dropout(p=dropout)))
                layers.append(('linear%d' % i, nn.Linear(input_size, size)))
                layers.append(('tanh%d' % i, nn.Tanh()))
                input_size = size
        if dropout > 0.:
          layers.append(('dropout', nn.Dropout(p=dropout)))
        layers.append(('linear', nn.Linear(input_size, output_size)))
        self.layer = nn.Sequential(OrderedDict(layers))
        self.activation = activation
        if not aggregator in ['sum', 'avg']:
            raise ValueError("I can only aggregate outputs using 'sum' or 'avg'")
        self.aggregator = aggregator

    def forward(self, x, mask, lengths):
        # [B, T, d]
        y = self.layer(x)
        if not self.activation is None:
            y = self.activation(y)
        # [B, d]
        s = (y * mask.unsqueeze(-1).float()).sum(dim=1)
        if self.aggregator == 'avg':
            s /= lengths.unsqueeze(-1).float()
        return y, s
```

```python
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

class LSTMEncoder(Encoder):
    """
    This module encodes a sequence using a bidirectional LSTM
     it returns the final state
     and the hidden states at each time step. Note: we concatenate representations
     from the two directions.
    """

    def __init__(self, in_features,
                 hidden_size: int = 200,
                 batch_first: bool = True,
                 bidirectional: bool = True):
        """
        :param in_features:
        :param hidden_size:
        :param batch_first:
        :param bidirectional:
        """
        pass

    def forward(self, x, mask, lengths):
        """
        Encode sentence x
        :param x: sequence of word embeddings, shape [B, T, I]
        :param mask: byte mask that is 0 for invalid positions, shape [B, T]
        :param lengths: the lengths of each input sequence [B]
        :return:
            outputs [B, T, O]
            final state [B, O]
        """
        pass
```

```python
# SOLUTION

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

class LSTMEncoder(Encoder):
    """
    This module encodes a sequence into a single vector using an LSTM,
     it also returns the hidden states at each time step.
    """

    def __init__(self, in_features, hidden_size: int = 200,
                 batch_first: bool = True,
                 bidirectional: bool = True):
        """
        :param in_features:
        :param hidden_size:
        :param batch_first:
        :param bidirectional:
        """
        super(LSTMEncoder, self).__init__()
        self.lstm = nn.LSTM(in_features, hidden_size, batch_first=batch_first,
                            bidirectional=bidirectional)

    def forward(self, x, mask, lengths):
        """
        Encode sentence x
        :param x: sequence of word embeddings, shape [B, T, E]
        :param mask: byte mask that is 0 for invalid positions, shape [B, T]
        :param lengths: the lengths of each input sequence [B]
        :return:
        """

        packed_sequence = pack_padded_sequence(x, lengths, batch_first=True)
        outputs, (hx, cx) = self.lstm(packed_sequence)
        outputs, _ = pad_packed_sequence(outputs, batch_first=True)

        # classify from concatenation of final states
        if self.lstm.bidirectional:
            final = torch.cat([hx[-2], hx[-1]], dim=-1)
        else:  # classify from final state
            final = hx[-1]

        return outputs, final
```

Here is some helper code to select and return an encoder:

```python
def get_encoder(layer, in_features, hidden_size, bidirectional=True):
    """Returns the requested layer."""

    if layer == "bow":
        return BagOfWordsEncoder()
    elif layer == 'ff':
        return FFEncoder(
            in_features,
            2 * hidden_size,   # for convenience
            hidden_sizes=[hidden_size],
            aggregator='avg')
    elif layer == "lstm":
        return LSTMEncoder(
            in_features,
            hidden_size,
            bidirectional=bidirectional)
    else:
        raise ValueError("Unknown layer")
```

```python

```

# Sentiment Classification with Latent Rationale

A latent rationale is a compact and informative fragment of the input based on which a NN classifier makes its decisions. [Lei et al (2016)](http://aclweb.org/anthology/D16-1011) proposed to induce such rationales along with a regression model for multi-aspect sentiment analsysis, their model is trained via REINFORCE on a dataset of beer reviews.

*Remark:* the model we will develop here can be seen as a probabilistic version of their model. The rest of this notebook focus on our own probabilitisc view of the model.

The picture below depicts our latent-variable model for rationale extraction:

where we augment the model with a collection of latent variables $z = \langle z_1, \ldots, z_n\rangle$ where $z_i$ is a binary latent variable. Each latent variable $z_i$ regulates whether or not the input $x_i$ is available to the classifier.  We use $x \odot z$ to denote the selected words, which, in the terminology of Lei et al, is a latent rationale.

Again the classifier parameterises a Categorical distribution over $K=5$ outcomes, though this time it can encode only a selection of the input:

\begin{align}
    Z_i & \sim \text{Bern}(p_1) \\
    Y|z,x &\sim \text{Cat}(f(x \odot z; \theta))
\end{align}

where we have a shared and fixed Bernoulli prior (with parameter $p_1$) for all $n$ latent variables.

Here is an example design for $f$:

\begin{align}
\mathbf x_i &= z_i \, \text{glove}(x_i) \\
\mathbf t_1^n, \mathbf h &= \text{encoder}(\mathbf x_1^n; \theta_{\text{enc}}) \\
f(x \odot z; \theta) &= \text{softmax}(\text{dense}_K(\mathbf h; \theta_{\text{output}}))
\end{align}

where:
* $z_i$ either leaves $\mathbf x_i$ unchanged or turns it into a vector of zeros;
* the encoder only sees features from selected inputs, i.e. $x_i$ for which $z_i = 1$;
* $\text{dense}_K$ is a linear layer with $K=5$ outputs.

## Prior

Our prior is a Bernoulli with fixed parameter $0 < p_1 < 1$:

\begin{align}
Z_i|p_1 & \sim \text{Bern}(p_1)
\end{align}

whose pmf is $\text{Bern}(z_i|p_1) = p_1^{z_i}\times (1-p_1)^{1-z_i}$.

As we will be using Bernoulli priors and posteriors, it is a good idea to implement a Bernoulli class:

```python
class Bernoulli:
    """
    This class encapsulates a collection of Bernoulli distributions.
    Each Bernoulli is uniquely specified by p_1, where
        Bernoulli(X=x|p_1) = pow(p_1, x) * pow(1 - p_1, 1 - x)
    is the Bernoulli probability mass function (pmf).
    """

    def __init__(self, logits=None, probs=None):
        """
        We can specify a Bernoulli distribution via a logit or a probability.
         You need to specify at least one, and if you specify both, beware that
         in this implementation logits will be used.

        Recall that: probs = sigmoid(logits).

        :param logits: a tensor of logits (a logit is defined as log (p_1/p_0))
            where p_0 = 1 - p_1
        :param probs: a tensor of probabilities, each in (0, 1)

        """
        pass

    def sample(self):
        """Returns a single sample with the same shape as the parameters"""
        pass

    def ste_sample(self):
        """
        Returns a single sample with the same shape as the parameters using the straight-through estimator.
        That is, you must make sure the forward pass produces a discrete sample
        while the backward pass differentiates the probability value (rather than the sample).

        **This is an extra, don't bother implementing it initially**
        """
        pass

    def log_pmf(self, x):
        """
        Assess the log probability of a sample.

        :param x: either a single sample (0 or 1) or a tensor of samples with the same shape as the parameters.
        :returns: tensor with log probabilities with the same shape as parameters
            (if the input is a single sample we broadcast it to the shape of the parameters)
        """
        pass

    def kl(self, other: 'Bernoulli'):
        """
        Compute the KL divergence between two Bernoulli distributions (from self to other).

        :return: KL[self||other] with same shape parameters
        """
        pass

```

```python
# SOLUTION
from torch.distributions import Bernoulli as PyTBernoulli
import torch.nn.functional as F

class Bernoulli:
    """
    This class encapsulates a collection of Bernoulli distributions.
    Each Bernoulli is uniquely specified by p_1, where
        Bernoulli(X=x|p_1) = pow(p_1, x) + pow(1 - p_1, 1 - x)
    is the Bernoulli probability mass function (pmf).
    """

    def __init__(self, logits=None, probs=None):
        """
        We can specify a Bernoulli distribution via a logit or a probability.
         You need to specify at least one, and if you specify both, beware that
         in this implementation logits will be used.

        Recall that: probs = sigmoid(logits).

        :param logits: a tensor of logits (a logit is defined as log (p_1/p_0))
            where p_0 = 1 - p_1
        :param probs: a tensor of probabilities, each in (0, 1)

        """
        if probs is None and logits is None:
            raise ValueError('I need probabilities or logits')
        if logits is None:
            eps = torch.finfo(probs.dtype).eps
            probs = probs.clamp(min=eps, max=1 - eps)
            self.probs = probs
            self.logits = torch.log(probs) - torch.log1p(-probs)
        else:
            self.logits = logits
            self.probs = torch.sigmoid(logits)
        # stable computation of log prob
        self.log_p1 = - F.softplus(-self.logits)
        self.log_p0 = - F.softplus(self.logits)

    def sample(self):
        """Returns a sample with the same shape as the parameters"""
        return (torch.rand_like(self.probs) < self.probs).float()

    def ste_sample(self):
        """
        Returns a single sample with the same shape as the parameters using the straight-through estimator.
        That is, you must make sure the forward pass produces a discrete sample
        while the backward pass differentiates the probability value (rather than the sample).

        **This is an extra, don't bother implementing it initially**
        """
        z = (torch.rand_like(self.probs) < self.probs).float()
        # Note how the forward pass evaluates to z
        #  and the backward pass evaluates to p'
        return (z - self.probs).detach() + self.probs - self.probs.detach()

    def log_pmf(self, x):
        """
        Assess the log probability of a sample.
        :param x: either a single sample (0 or 1) or a tensor of samples with the same shape as the parameters.
        :returns: tensor with log probabilities with the same shape as parameters
            (if the input is a single sample we broadcast it to the shape of the parameters)
        """
        # x * torch.log(self.probs) + (1 - x) * torch.log(1. - self.probs)
        return torch.where(x == 1., self.log_p1, self.log_p0)

    def kl(self, other: 'Bernoulli'):
        """
        Compute the KL divergence between two Bernoulli distributions (from self to other).

        :return: KL[self||other] with same shape parameters
        """
        return self.probs * (self.log_p1 - other.log_p1) + (1 - self.probs) * (self.log_p0 - other.log_p0)
```

```python
def test_bernoulli_shapes():
    batch_size = 2
    latent_dim = 3
    p = Bernoulli(torch.sigmoid(torch.rand([batch_size, latent_dim])))
    q = Bernoulli(torch.sigmoid(torch.rand([batch_size, latent_dim])))
    assert p.log_pmf(torch.ones([batch_size, latent_dim]).long()).size() == torch.Size([batch_size, latent_dim])
    assert p.kl(q).size() == torch.Size([batch_size, latent_dim])

test_bernoulli_shapes()
```

```python

def test_bernoulli_sampling():
    import numpy as np

    # create a simple Bernoulli
    p = Bernoulli(probs=torch.full([1], 0.2))
    assert p.sample().numpy().shape == (1,), 'Wrong sample shape'
    np.testing.assert_almost_equal(
        np.mean([p.sample().numpy() for _ in range(1000)]),
        0.2,
        decimal=1,
        err_msg="Bad sample mean")

    # create a batch of 1000 Bernoulli distributions
    p = Bernoulli(probs=torch.full([1000, 1], 0.2))
    assert p.sample().numpy().shape == (1000,1), 'Wrong sample shape'
    np.testing.assert_almost_equal(
        np.mean(p.sample().numpy()),
        0.2,
        decimal=1,
        err_msg="Bad sample mean")

    # create a batch of 1000 Bernoulli distributions
    p = Bernoulli(logits=torch.full([1000, 1], 0.3).log())
    assert p.sample().numpy().shape == (1000,1), 'Wrong sample shape'
    np.testing.assert_almost_equal(
        np.mean(p.sample().numpy()),
        0.2,
        decimal=1,
        err_msg="Bad sample mean")

test_bernoulli_sampling()
```

## Classifier

The classifier encodes only a selection of the input, which we denote $x \odot z$, and parameterises a Categorical distribution over $5$ outcomes (sentiment levels).

Thus let's implement a Categorical distribution (we will only need to be able to assess its lgo pmf):

```python
class Categorical:

    def __init__(self, log_probs):
        # [B, K]: class probs
        self.log_probs = log_probs

    def log_pmf(self, x):
        """
        :param x: [B] integers (targets)
        :returns: [B] scalars (log probabilities)
        """
        pass

    def argmax(self):
        """
        Return the argmax prediction
        :returns: [B] class ids
        """
        pass
```

```python
#SOLUTION
class Categorical:

    def __init__(self, log_probs):
        # [B, K]: class probs
        self.log_probs = log_probs

    def log_pmf(self, x):
        """
        :param x: [B] integers
        """
        # [B, 1]
        log_p = torch.gather(self.log_probs, 1, x.unsqueeze(-1))
        # [B]
        return log_p.squeeze(-1)

    def argmax(self):
        """
        Return the argmax prediction
        :returns: [B] class ids
        """
        return self.log_probs.argmax(-1)
```

```python
def test_categorical_shapes():
    batch_size = 2
    nb_classes = 5
    p = Categorical(F.softmax(torch.rand([batch_size, nb_classes]), -1))
    assert p.log_pmf(torch.ones(batch_size).long()).size() == torch.Size([batch_size]), "Did you squeeze the category dimension?"
    assert p.argmax().size() == torch.Size([batch_size]), "Did you squeeze the category dimension?"
test_categorical_shapes()
```

and a classifier architecture:

* implement the forward method

```python
class Classifier(nn.Module):
    """
    The Encoder takes an input text (and rationale z) and computes p(y|x,z)
    """

    def __init__(self,
                 embed:        nn.Embedding = None,
                 hidden_size:  int = 200,
                 output_size:  int = 1,
                 dropout:      float = 0.1,
                 layer:        str = "pass",
                 ):

        super(Classifier, self).__init__()

        emb_size = embed.weight.shape[1]
        enc_size = hidden_size * 2
        # Here we embed the words
        self.embed_layer = nn.Sequential(
            embed
            # , nn.Dropout(p=dropout)
        )

        self.enc_layer = get_encoder(layer, emb_size, hidden_size)

        # and here we predict categorical parameters
        self.output_layer = nn.Sequential(
            nn.Dropout(p=dropout),
            nn.Linear(enc_size, output_size),
            nn.LogSoftmax(dim=-1)
        )

        self.report_params()

    def report_params(self):
        count = 0
        for name, p in self.named_parameters():
            if p.requires_grad and "embed" not in name:
                count += np.prod(list(p.shape))
        print("{} #params: {}".format(self.__class__.__name__, count))

    def forward(self, x, mask, z) -> Categorical:
        """
        :params x: [B, T, I] word representations
        :params mask: [B, T] indicates valid positions
        :params z: [B, T] binary selectors
        :returns: one Categorical distribution per instance in the batch
          each conditioning only on x_i for which z_i = 1
        """
        pass
```

```python
#SOLUTION
class Classifier(nn.Module):
    """
    The Encoder takes an input text (and rationale z) and computes p(y|x,z)
    """

    def __init__(self,
                 embed:        nn.Embedding = None,
                 hidden_size:  int = 200,
                 output_size:  int = 5,
                 dropout:      float = 0.1,
                 layer:        str = "pass",
                 ):

        super(Classifier, self).__init__()

        emb_size = embed.weight.shape[1]
        enc_size = hidden_size * 2
        # Here we embed the words
        self.embed_layer = nn.Sequential(
            embed
        )

        self.enc_layer = get_encoder(layer, emb_size, hidden_size)

        # and here we predict categorical parameters
        self.output_layer = nn.Sequential(
            nn.Dropout(p=dropout),
            nn.Linear(enc_size, output_size),
            nn.LogSoftmax(dim=-1)
        )

        self.report_params()

    def report_params(self):
        count = 0
        for name, p in self.named_parameters():
            if p.requires_grad:
                count += np.prod(list(p.shape))
        print("{} #params: {}".format(self.__class__.__name__, count))

    def forward(self, x, mask, z) -> Categorical:
        """
        :params x: [B, T, I] word representations
        :params mask: [B, T] indicates valid positions
        :params z: [B, T] binary selectors
        :returns: one Categorical distribution per instance in the batch
          each conditioning only on x_i for which z_i = 1
        """

        rnn_mask = mask
        emb = self.embed_layer(x)

        # [B, T]
        rnn_mask = z > 0.
        # [B, T, 1]
        z_mask = z.unsqueeze(-1).float()
        # [B, T, E]
        emb = emb * z_mask

        lengths = mask.long().sum(1)

        # encode the sentence
        _, final = self.enc_layer(emb, rnn_mask, lengths)

        # predict sentiment from final state(s)
        log_probs = self.output_layer(final)
        return Categorical(log_probs)
```

## Inference

Computing the log-likelihood of an observation requires marginalising over assignments of $z$:

\begin{align}
P(y|x,\theta,p_1) &= \sum_{z_1 = 0}^1 \cdots \sum_{z_n=0}^1 P(z|p_1)\times P(y|x,z, \theta) \\
&= \sum_{z_1 = 0}^1 \cdots \sum_{z_n=0}^1 \left( \prod_{i=1}^n \text{Bern}(z_i|p_1)\right) \times \text{Cat}(y|f(x \odot z; \theta))
\end{align}

This is clearly intractable: there are $2^n$ possible assignments to $z$ and because the classifier conditions on all latent selectors, there's no way to simplify the expression.

We will avoid computing this intractable marginal by instead employing an independently parameterised inference model.
This inference model $Q(z|x, y, \lambda)$ is an approximation to the true postrerior $P(z|x, y, \theta, p_1)$, and we use $\lambda$ to denote its parameters.

We make a *mean field* assumption, whereby we model latent variables independently given the input:
\begin{align}
Q(z|x, y, \lambda)
    &= \prod_{i=1}^{n} Q(z_i|x; \lambda) \\
    &= \prod_{i=1}^{n} \text{Bern}(z_i|g_i(x; \lambda))
\end{align}

where $g(x; \lambda)$ is a NN that maps from $x = \langle x_1, \ldots, x_n\rangle$ to $n$ Bernoulli parameters, each of which, is a probability value (thus $0 < g_i(x; \lambda) < 1$).

Note that though we could condition on $y$ for approximate posterior inference, we are opportunistically leaving it out. This way, $Q$ is directly available at test time for making predictions. The figure below is a graphical depiction of the inference model (we show a dashed arrow from $y$ to $z$ to remind you that in principle the label is also available).

Here is an example design for $g$:
\begin{align}
\mathbf x_i &= \text{glove}(x_i) \\
\mathbf t_1^n, \mathbf h &= \text{encoder}(\mathbf x_1^n; \lambda_{\text{enc}}) \\
g_i(x; \lambda) &= \sigma(\text{dense}_1(\mathbf t_i; \lambda_{\text{output}}))
\end{align}
where
* $\text{glove}$ is a pre-trained embedding function;
* $\text{dense}_1$ is a dense layer with a single output;
* and $\sigma(\cdot)$ is the sigmoid function, necessary to parameterise a Bernoulli distribution.

From now on we will write $Q(z|x, \lambda)$, that is, without $y$.

Here we implement this product of Bernoulli distributions:

* implement $g$ in the constructor
* and the forward pass

```python
class ProductOfBernoullis(nn.Module):
    """
    This is an inference network that parameterises independent Bernoulli distributions.
    """

    def __init__(self,
                 embed:       nn.Embedding,
                 hidden_size: int = 200,
                 layer:       str = "bow"
                 ):
        """
        :param embed: an embedding layer
        :param hidden_suze: hidden size for transformed inputs
        :param layer: 'bow' for BoW encoding
          you may alternatively implement and 'lstm' option
          which uses a biLSTM to transform the inputs
        """
        super(ProductOfBernoullis, self).__init__()
        # 1. we should have an embedding layer
        # 2. we may transform the representations
        # 3. and we should compute parameters for Bernoulli distributions
        pass

    def report_params(self):
        count = 0
        for name, p in self.named_parameters():
            if p.requires_grad and "embed" not in name:
                count += np.prod(list(p.shape))
        print("{} #params: {}".format(self.__class__.__name__, count))

    def forward(self, x, mask) -> Bernoulli:
        """
        It takes a tensor of tokens (integers)
         and predicts a Bernoulli distribution for each position.

        :param x: [B, T]
        :param mask: [B, T]
        :returns: Bernoulli
        """
        pass
```

```python
#SOLUTION
class ProductOfBernoullis(nn.Module):
    """
    This is an inference network that parameterises independent Bernoulli distributions.
    """

    def __init__(self,
                 embed:       nn.Embedding,
                 hidden_size: int = 200,
                 layer:       str = "bow"
                 ):
        """
        :param embed: an embedding layer
        :param hidden_suze: hidden size for transformed inputs
        :param layer: 'bow' for BoW encoding
          you may alternatively implement and 'lstm' option
          which uses a biLSTM to transform the inputs
        """

        super(ProductOfBernoullis, self).__init__()

        emb_size = embed.weight.shape[1]
        # I will use twice the units
        #  just to make the output as large as that of a biLSTM
        enc_size = hidden_size * 2

        self.embed_layer = nn.Sequential(embed)
        self.enc_layer = get_encoder(layer, emb_size, hidden_size)
        self.logit_layer = nn.Linear(enc_size, 1, bias=True)

        self.report_params()

    def report_params(self):
        count = 0
        for name, p in self.named_parameters():
            if p.requires_grad:
                count += np.prod(list(p.shape))
        print("{} #params: {}".format(self.__class__.__name__, count))

    def forward(self, x, mask) -> Bernoulli:
        """
        It takes a tensor of tokens (integers)
         and predicts a Bernoulli distribution for each position.

        :param x: [B, T]
        :param mask: [B, T]
        :returns: Bernoulli
        """

        # encode sentence
        # [B]
        lengths = mask.long().sum(1)
        # [B, T, E]
        emb = self.embed_layer(x)
        # [B, T, d]
        h, _ = self.enc_layer(emb, mask, lengths)

        # compute parameters for Bernoulli p(z|x)
        # [B, T, 1] Bernoulli distributions
        logits = self.logit_layer(h)
        # [B, T]
        logits = logits.squeeze(-1)
        return Bernoulli(logits=logits)
```

## Parameter Estimation

In variational inference, our objective is to maximise the *evidence lowerbound* (ELBO):

\begin{align}
\log P(y|x) &\ge \mathbb E_{Q(z|x, y, \lambda)}\left[ \log P(y|x, z, \theta, p_1) \right] - \text{KL}(Q(z|x, y, \lambda) || P(z|p_1)) \\
\text{ELBO}&\overset{\text{MF}}{=}\mathbb E_{Q(z|x, y, \lambda)}\left[ \log P(y|x, z, \theta, p_1) \right] - \sum_{i=1}^n \text{KL}(Q(z_i|x, \lambda) || P(z_i|p_1))
\end{align}

where the *mean field* assumption we made implies that the KL term is simply a sum of KL divergences from a Bernoulli posterior to a Bernoulli prior.

Note that the ELBO remains intractable, namely, solving the expectation in closed form still requires $2^n$ evaluations of the classifier network. Though unlike the true posterior $P(z|x,y, \lambda)$, the approximation $Q(z|x,\lambda)$ is tractable (it does not require an intractable normalisation) and can be used to obtain gradient estimates based on samples.

### Gradient of the classifier network

For the classifier, we encounter no problem:

\begin{align}
\nabla_\theta \text{ELBO} &=\nabla_\theta\sum_{z} Q(z|x, \lambda)\log P(y|x,z,\theta) - \underbrace{\nabla_\theta \sum_{i=1}^n \text{KL}(Q(z_i|x, \lambda) || P(z_i|p_1))}_{\color{blue}{0}}  \\
&=\sum_{z} Q(z|x, \lambda)\nabla_\theta\log P(y|x,z,\theta) \\
&= \mathbb E_{Q(z|x, \lambda)}\left[\nabla_\theta\log P(y|x,z,\theta) \right] \\
&\overset{\text{MC}}{\approx} \frac{1}{S} \sum_{s=1}^S \nabla_\theta \log P(y|x, z^{(s)}, \theta)
\end{align}
where $z^{(s)} \sim Q(z|x,\lambda)$.

### Gradient of the inference network

For the inference model, we have to use the *score function estimator* (a.k.a. REINFORCE):

\begin{align}
\nabla_\lambda \text{ELBO} &=\nabla_\lambda\sum_{z} Q(z|x, \lambda)\log P(y|x,z,\theta) - \nabla_\lambda \underbrace{\sum_{i=1}^n \text{KL}(Q(z_i|x, \lambda) || P(z_i|p_1))}_{ \color{blue}{\text{tractable} }}  \\
&=\sum_{z} \nabla_\lambda Q(z|x, \lambda)\log P(y|x,z,\theta) - \sum_{i=1}^n \nabla_\lambda \text{KL}(Q(z_i|x, \lambda) || P(z_i|p_1))   \\
&=\sum_{z}  \underbrace{Q(z|x, \lambda) \nabla_\lambda \log Q(z|x, \lambda)}_{\nabla_\lambda Q(z|x, \lambda)} \log P(y|x,z,\theta) - \sum_{i=1}^n \nabla_\lambda \text{KL}(Q(z_i|x, \lambda) || P(z_i|p_1))   \\
&= \mathbb E_{Q(z|x, \lambda)}\left[ \log P(y|x,z,\theta) \nabla_\lambda \log Q(z|x, \lambda) \right] - \sum_{i=1}^n \nabla_\lambda \text{KL}(Q(z_i|x, \lambda) || P(z_i|p_1))   \\
&\overset{\text{MC}}{\approx} \left(\frac{1}{S} \sum_{s=1}^S  \log P(y|x, z^{(s)}, \theta) \nabla_\lambda \log Q(z^{(s)}|x, \lambda)  \right) - \sum_{i=1}^n \nabla_\lambda \text{KL}(Q(z_i|x, \lambda) || P(z_i|p_1))
\end{align}

where $z^{(s)} \sim Q(z|x,\lambda)$.

## Implementation

Let's implement the model and the loss (negative ELBO). We work with the notion of a *surrogate loss*, that is, a computation node whose gradients wrt to parameters are equivalent to the gradients we need.

For a given sample $z \sim Q(z|x, \lambda)$, the following is a single-sample surrogate loss:

\begin{align}
\mathcal S(\theta, \lambda|x, y) = \log P(y|x, z, \theta) + \color{red}{\text{detach}(\log P(y|x, z, \theta) )}\log Q(z|x, \lambda) - \sum_{i=1}^n \text{KL}(Q(z_i|x, \lambda) || P(z_i|\phi))
\end{align}
where we introduce an auxiliary function such that
\begin{align}
\text{detach}(f(\alpha))  &= h(\alpha) \\
\nabla_\beta \text{detach}(h(\alpha))  &= 0
\end{align}
or in words, *detach* does not alter the forward call of its argument function $h$, but it alters $h$'s backward call by setting gradients to zero.

Show that it's gradients wrt $\theta$ and $\lambda$ are exactly what we need:

\begin{align}
\nabla_\theta \mathcal S(\theta, \lambda|x, y) = \color{red}{?}
\end{align}

\begin{align}
\nabla_\lambda \mathcal S(\theta, \lambda|x, y) = \color{red}{?}
\end{align}

**Solution**

\begin{align}
\nabla_\theta \mathcal S(\theta, \lambda|x, y) = \nabla_\theta \log P(y|x, z, \theta) + 0
\end{align}

\begin{align}
\nabla_\lambda \mathcal S(\theta, \lambda|x, y) &= 0 + \underbrace{\log Q(z|x, \lambda)\nabla_\lambda \log P(y|x, z, \theta)  + \log P(y|x, z, \theta) \nabla_\lambda \log Q(z|x, \lambda)}_{\text{chain rule}} \\
&= 0+ 0 + \log P(y|x, z, \theta) \nabla_\lambda \log Q(z|x, \lambda)
\end{align}

Implement the forward pass and loss below:

```python
class VAE(nn.Module):
    """

    Classifier model:
        Z_i ~ Bern(p_1) for i in 1..n
        Y|x,z ~ Cat(f([x_i if z_i 1 else 0 for i in 1..n ]))

    Inference model:
        Z_i|x ~ Bern(b_i) for i in 1..n
            where b_i = g_i(x)

    Objective:
        Single-sample MC estimate of ELBO

    Loss:
        Surrogate loss

    Consists of:
        - a product of Bernoulli distributions inference network
        - a classifier network
    """

    def __init__(self,
                 vocab:       object = None,
                 vocab_size:  int = 0,
                 emb_size:    int = 200,
                 hidden_size: int = 200,
                 num_classes: int = 5,
                 prior_p1:    float = 0.3,
                 dropout:     float = 0.1,
                 layer_cls:   str = 'bow',
                 layer_inf:   str = 'bow',
                 ):
        """
        :param vocab: Vocabulary
        :param vocab_size: necessary for embedding layer
        :param emb_size: dimensionality of embedding layer
        :param hidden_size: dimensionality of hidden layers
        :param num_classes: number of classes
        :param prior_p1: (scalar) prior Bernoulli parameter
        :param dropout: (scalar) dropout rate
        :param layer_cls: type of encoder for classification
        :param layer_inf: type of encoder for inference
        """
        super(VAE, self).__init__()

        self.vocab = vocab
        self.embed = embed = nn.Embedding(vocab_size, emb_size, padding_idx=1)

        self.cls_net = Classifier(
            embed=embed,
            hidden_size=hidden_size,
            output_size=num_classes,
            dropout=dropout,
            layer=layer_cls)

        self.inference_net = ProductOfBernoullis(
            embed=embed,
            hidden_size=hidden_size,
            layer=layer_inf)

        self._prior_p1 = prior_p1

    def generative_parameters(self):
        return self.cls_net.parameters()

    def inference_parameters(self):
        return self.inference_net.parameters()

    def get_prior(self, shape, device, dtype=torch.float32) -> Bernoulli:
        """
        Return a collection of independent prior Bernoulli distributions with a given shape.

        :param shape: typically a pair (batch_size, max_time)
        :param device: a torch device
        :returns: Bernoulli object whose probs have the given shape
            and the value is prior_p1
        """
        return Bernoulli(probs=torch.full(shape, self._prior_p1, device=device, dtype=dtype))

    def get_posterior(self, x, mask) -> Bernoulli:
        """
        Infer a collection of independent Bernoulli posteriors
         with shape [B, T]
        """
        return self.inference_net(x, mask)

    def get_likelihood(self, x, mask, z) -> Categorical:
        """
        Generate a sequence z with inference model,
         then predict with rationale xz, that is, x masked by z.

        :param x: [B, T] documents
        :return:
            Categorical distributions P(y|x, z)
            Bernoulli distributions Q(z|x)
            Single sample z ~ Q(z|x) used for the conditional P(y|x, z)
        """
        py = self.cls_net(x, mask, z)
        return py

    def forward(self, *args, **kwargs):
        """
        See get_loss
        """
        return self.get_loss(args, kwargs)

    def get_loss(self, x, mask, y,
                 iter_i=0,
                 kl_weight=1.0,
                 min_kl=0.0,
                 ll_mean=0.,
                 ll_std=1.,
                 **kwargs):
        """
        This computes the loss for the whole model.

        :param x: inputs [B, T]
        :param mask: identifies valid positions [B, T]
        :param y: target labels [B]
        :param iter_i: indicates the iteration
        :param kl_weight: (scalar) multiplies the KL term
        :param min_kl: (scalar) sets a minimum for the KL (aka free bits)
        :param ll_mean: (scalar) running average of reward
        :param ll_std: (scalar) running standard deviation of reward
        :return: loss (torch scalar node), terms (dict)

            terms is a dict that holds the scalar items involved in the loss
            e.g. `terms['ll'] = ll.item()` is the log-likelihood term

            You should log the following:

            - the ELBO: terms['elbo']
            - the likelihood term of the ELBO: terms['ll']
            - the kl before clipping and annealing: terms['kl']
            - the kl after clipping and annealing: terms['kl_']
            - the learning signal in REINFORCE: terms['reward']
            - the standard deviation of the learning signal: terms['ll_std']
            - the ratio of selected words (average z): terms['selected']

            Consider tracking the following:
            Single-sample ELBO: terms['elbo']
            Log-Likelihood log P(y|x,z): terms['ll']
            KL: terms['kl']
            Score function surrogate log P(y|z, x) log Q(z|x): terms['sf']
            Rate of selected words: terms['selected']
        """

        raise NotImplementedError("Implement me!")

        # Make complete surrogate loss (for backward)
        # [B]
        loss = None

        # Store terms worth tracking
        terms['elbo'] = 0.
        terms['kl_'] = 0.
        terms['kl'] = 0.
        terms['reward'] = 0.
        terms['ll'] = 0.
        terms['ll_std'] = 0.
        terms['selected'] = 0.

        # return loss for backward and terms worth tracking
        return loss.mean(), terms

    def get_ste_loss(self, x, mask, y,
                 iter_i=0,
                 # you may ignore the rest of the arguments for the time being
                 #  leave them as they are
                 kl_weight=1.0,
                 min_kl=0.0,
                 ll_mean=0.,
                 ll_std=1.,
                 **kwargs):
        """
        This is an extra:
            - use the (biased) straight-through estimator instead of the score function estimator
        """
        pass
```

```python
#SOLUTION
class VAE(nn.Module):
    """

    Classifier model:
        Z_i ~ Bern(p_1) for i in 1..n
        Y|x,z ~ Cat(f([x_i if z_i 1 else 0 for i in 1..n ]))

    Inference model:
        Z_i|x ~ Bern(b_i) for i in 1..n
            where b_i = g_i(x)

    Objective:
        Single-sample MC estimate of ELBO

    Loss:
        Surrogate loss

    Consists of:
        - a product of Bernoulli distributions inference network
        - a classifier network
    """

    def __init__(self,
                 vocab:       object = None,
                 vocab_size:  int = 0,
                 emb_size:    int = 200,
                 hidden_size: int = 200,
                 num_classes: int = 5,
                 prior_p1:    float = 0.3,
                 dropout:     float = 0.1,
                 layer_cls:   str = 'bow',
                 layer_inf:   str = 'bow',
                 ):
        """
        :param vocab: Vocabulary
        :param vocab_size: necessary for embedding layer
        :param emb_size: dimensionality of embedding layer
        :param hidden_size: dimensionality of hidden layers
        :param num_classes: number of classes
        :param prior_p1: (scalar) prior Bernoulli parameter
        :param dropout: (scalar) dropout rate
        :param layer_cls: type of encoder for classification
        :param layer_inf: type of encoder for inference
        """
        super(VAE, self).__init__()

        self.vocab = vocab
        self.embed = embed = nn.Embedding(vocab_size, emb_size, padding_idx=1)

        self.cls_net = Classifier(
            embed=embed,
            hidden_size=hidden_size,
            output_size=num_classes,
            dropout=dropout,
            layer=layer_cls)

        self.inference_net = ProductOfBernoullis(
            embed=embed,
            hidden_size=hidden_size,
            layer=layer_inf)

        self._prior_p1 = prior_p1

    def generative_parameters(self):
        return self.cls_net.parameters()

    def inference_parameters(self):
        return self.inference_net.parameters()

    def get_prior(self, shape, device, dtype=torch.float32) -> Bernoulli:
        """
        Return a collection of independent prior Bernoulli distributions with a given shape.

        :param shape: typically a pair (batch_size, max_time)
        :param device: a torch device
        :returns: Bernoulli object whose probs have the given shape
            and the value is prior_p1
        """
        return Bernoulli(probs=torch.full(shape, self._prior_p1, device=device, dtype=dtype))

    def get_posterior(self, x, mask) -> Bernoulli:
        """
        Infer a collection of independent Bernoulli posteriors
         with shape [B, T]
        """
        return self.inference_net(x, mask)

    def get_likelihood(self, x, mask, z) -> Categorical:
        """
        Generate a sequence z with inference model,
         then predict with rationale xz, that is, x masked by z.

        :param x: [B, T] documents
        :return:
            Categorical distributions P(y|x, z)
            Bernoulli distributions Q(z|x)
            Single sample z ~ Q(z|x) used for the conditional P(y|x, z)
        """
        py = self.cls_net(x, mask, z)
        return py

    def forward(self, *args, **kwargs):
        """
        Generate a sequence z with inference model,
         then predict with rationale xz, that is, x masked by z.

        :param x: [B, T] documents
        :return:
            Categorical distributions P(y|x, z)
            Bernoulli distributions Q(z|x)
            Single sample z ~ Q(z|x) used for the conditional P(y|x, z)
        """
        return self.get_loss(args, kwargs)

    def get_loss(self, x, mask, y,
                 iter_i=0,
                 # you may ignore the rest of the arguments for the time being
                 #  leave them as they are
                 kl_weight=1.0,
                 min_kl=0.0,
                 ll_mean=0.,
                 ll_std=1.,
                 **kwargs):
        """
        This computes the loss for the whole model.

        :param x: inputs [B, T]
        :param mask: identifies valid positions [B, T]
        :param y: target labels [B]
        :param iter_i: indicates the iteration
        :param kl_weight: (scalar) multiplies the KL term
        :param min_kl: (scalar) sets a minimum for the KL (aka free bits)
        :param ll_mean: (scalar) running average of reward
        :param ll_std: (scalar) running standard deviation of reward
        :return: loss (torch scalar node), terms (dict)

            terms is a dict that holds the scalar items involved in the loss
            e.g. `terms['ll'] = ll.item()` is the log-likelihood term

            You should log the following:

            - the ELBO: terms['elbo']
            - the likelihood term of the ELBO: terms['ll']
            - the kl before clipping and annealing: terms['kl']
            - the kl after clipping and annealing: terms['kl_']
            - the learning signal in REINFORCE: terms['reward']
            - the standard deviation of the learning signal: terms['ll_std']
            - the ratio of selected words (average z): terms['selected']

            Consider tracking the following:
            Single-sample ELBO: terms['elbo']
            Log-Likelihood log P(y|x,z): terms['ll']
            KL: terms['kl']
            Score function surrogate log P(y|z, x) log Q(z|x): terms['sf']
            Rate of selected words: terms['selected']
        """

        lengths = mask.sum(1).float()
        batch_size = mask.size(0)
        terms = dict()

        # Infer a posterior distribution Z|x
        qz = self.get_posterior(x, mask)

        # Sample latent variables Z
        # [B, T]
        z = qz.sample()
        # Apply input mask
        z = torch.where(mask, z, torch.zeros_like(z))

        # Obtain a conditional distribution Y|z
        # [B, T, K]
        py = self.get_likelihood(x, mask, z)

        # Log-likelihood log P(y|x,z)
        # [B]
        ll = py.log_pmf(y)

        # KL(q||p)
        # [B, T]
        pz = self.get_prior(z.size(), device=z.device, dtype=z.dtype)

        # Compute KL divergence
        # [B, T]
        kl = qz.kl(pz)
        # Mask invalid input positions
        # [B, T]
        kl = torch.where(mask, kl, torch.zeros_like(kl))

        # Compute the log density of the sample
        # [B, T]
        log_q_z = qz.log_pmf(z)
        # and mask invalid positions
        log_q_z = torch.where(mask, log_q_z, torch.zeros_like(log_q_z))
        # We have independent Bernoullis, thus we just sum their log probabilities
        # [B]
        log_q_z = log_q_z.sum(1)

        # surrogate objective for score function estimator
        # [B]
        reward = (ll - ll_mean) / ll_std
        sf_surrogate = reward.detach() * log_q_z

        # [B]
        kl = kl.sum(dim=-1)
        # KL may require annealing and free-bits
        kl_ = torch.max(torch.full_like(kl, min_kl), kl) * kl_weight

        # Make complete surrogate loss (for backward)
        # [B]
        loss = - (ll + sf_surrogate - kl_)

        # Store terms worth tracking
        terms['elbo'] = (ll - kl).mean().item()
        terms['kl_'] = kl_.mean().item()
        terms['kl'] = kl.mean().item()
        terms['reward'] = reward.mean().item()
        terms['ll'] = ll.mean().item()
        terms['ll_std'] = ll.std().item()
        terms['selected'] = (z.sum(1) / lengths).mean().item()

        # return loss for backward and terms worth tracking
        return loss.mean(), terms

    def get_ste_loss(self, x, mask, y,
                 iter_i=0,
                 # you may ignore the rest of the arguments for the time being
                 #  leave them as they are
                 kl_weight=1.0,
                 min_kl=0.0,
                 ll_mean=0.,
                 ll_std=1.,
                 **kwargs):
        """Here we use the biased straight-through estimator instead of the score function estimator"""

        lengths = mask.sum(1).float()
        batch_size = mask.size(0)
        terms = dict()

        # Infer a posterior distribution Z|x
        qz = self.get_posterior(x, mask)

        # Sample latent variables Z
        # [B, T]
        z = qz.ste_sample()
        # Apply input mask
        z = torch.where(mask, z, torch.zeros_like(z))

        # Obtain a conditional distribution Y|z
        # [B, T, K]
        py = self.get_likelihood(x, mask, z)

        # Log-likelihood log P(y|x,z)
        # [B]
        ll = py.log_pmf(y)

        # KL(q||p)
        # [B, T]
        pz = self.get_prior(z.size(), device=z.device, dtype=z.dtype)

        # Compute KL divergence
        # [B, T]
        kl = qz.kl(pz)
        # Mask invalid input positions
        # [B, T]
        kl = torch.where(mask, kl, torch.zeros_like(kl))

        # [B]
        kl = kl.sum(dim=-1)
        # KL may require annealing and free-bits
        kl_ = torch.max(torch.full_like(kl, min_kl), kl) * kl_weight

        # Make complete surrogate loss (for backward)
        # [B]
        loss = - (ll - kl_)

        # Store terms worth tracking
        terms['elbo'] = (ll - kl).mean().item()
        terms['kl_'] = kl_.mean().item()
        terms['kl'] = kl.mean().item()
        terms['ll'] = ll.mean().item()
        terms['ll_std'] = ll.std().item()
        terms['selected'] = (z.sum(1) / lengths).mean().item()
        terms['reward'] = 0.

        # return loss for backward and terms worth tracking
        return loss.mean(), terms
```

# Training loop

For you to focus on the machine learning, we provide you with functions to deal with mini-batching, evaluation, and training.

```python
# some helper code for mini batching
#  this will take care of annoying things such as
#  sorting training instances by length (necessary for pytorch's LSTM, for example)
from util import get_minibatch, prepare_minibatch, print_parameters
```

An intrinsic convergence criterion is validation ELBO, thus we help you compute it here:

```python
def eval_elbo(model: VAE, data, batch_size=25, device=None, iter_i=0, nb_samples=10):
    """
    Accuracy of a model on given data set (using minibatches)

    :return: validation ELBO, validation KL
    """

    model.eval()  # disable dropout
    total_elbo, total_kl, data_size = 0., 0., 0.
    for mb in get_minibatch(data, batch_size=batch_size, shuffle=False):
        x, y, _ = prepare_minibatch(mb, model.vocab, device=device)
        mask = (x != 1)
        data_size = data_size + x.size(0)
        with torch.no_grad():
            # Infer Z|x
            qz = model.get_posterior(x, mask)
            # [B, T]
            z_s = qz.sample()
            # Prior Z
            # [B, T]
            pz = model.get_prior(z_s.size(), device=z_s.device, dtype=z_s.dtype)
            # 1. KL (closed form)
            # [B]
            kl = torch.where(mask, qz.kl(pz), torch.zeros_like(z_s)).sum(-1)
            # 2. E_q[log P(y|x,z)]
            terms = []
            ll = 0.
            for s in range(nb_samples):
                # sth sample
                # [B, T]
                z_s = qz.sample()
                # Apply input mask
                z_s = torch.where(mask, z_s, torch.zeros_like(z_s))
                # Compute Y|x,z_s
                # [B, K]
                py = model.get_likelihood(x, mask, z_s)
                # [B]
                ll = ll + py.log_pmf(y)
            total_elbo = total_elbo + (ll / nb_samples - kl ).sum().item()
            total_kl = total_kl + kl.sum().item()
    return total_elbo / data_size, total_kl / data_size
```

An extrinsic convergence crition is classification accuracy, thus we help you compute it here.
We also make visualisations that highlight selected words (rationales).

```python
from collections import OrderedDict, defaultdict
import numpy as np

def decorate_token(t, z, hiding=True):
    """
    Either hide text that hasn't been selected, or highlight text that has been selected.
    """
    if hiding:
        return t if z == 1 else "_"
    else:
        dec = "**" if z == 1 else ""
        return dec + t + dec

def evaluate(model: VAE, data, batch_size=25, device=None, iter_i=0, nb_samples=10):
    """
    Accuracy of a model on given data set (using minibatches)
    :return: accuracy, average selection rate, list of rationales (each a string)
    """

    model.eval()  # disable dropout
    rationales = []
    nb_correct, data_size, selection_rate = 0., 0., 0.
    for mb in get_minibatch(data, batch_size=batch_size, shuffle=False):
        x, targets, reverse_map = prepare_minibatch(mb, model.vocab, device=device)
        mask = (x != 1)
        with torch.no_grad():

            # Infere Z|x
            qz = model.get_posterior(x, mask)

            # Deterministic predictions (via greedy argmax)
            # [B, T]
            z_max = (qz.probs >= 0.5).float()
            # mask invalid positions
            z_max = torch.where(mask, z_max, torch.zeros_like(z_max))
            # Compute Y|x,z
            # [B, K]
            py = model.get_likelihood(x, mask, z_max)
            # argmax class
            # [B]
            predictions = py.argmax()

            # Accuracy and selection rate
            nb_correct = nb_correct + (predictions == targets.view(-1)).sum().item()
            data_size = data_size + x.size(0)
            selection_rate = selection_rate + (z_max.sum(-1) / mask.sum(-1).float()).sum().item()

            # String rationales
            # reverse sort
            z_max = z_max.cpu().numpy()
            z_max = z_max[reverse_map]
            for idx in range(x.size(0)):  # iterate over instances in a mini batch
                example = []
                for ti, zi in zip(mb[idx].tokens, z_max[idx]):  # iterate over tokens in an instance
                    example.append(decorate_token(ti, zi))
                rationales.append((example, predictions[idx]))

    return nb_correct / data_size, selection_rate / data_size, rationales
```

Here we configure our training procedure. As there are several hyperparameters, let's us discuss the main ones briefly.

**Data**
* `training_path` where to find the training text file
* `dev_path` where to find the validation text file
* `test_path` where to find the test text file
* `word_vectors` path to pre-trained 300-dimensional English word vectors
* `subphrases` should we use annotation at subphrase level
* `min_phrase_length` shortest phrase allowed (when subphrases are used)
* `lowercase` should we lowercase the data

**Model**
* `prior_p1` this is the Bernoulli prior parameter $p_1$

**Architecture**
* `fix_emb` should we keep embeddings fixed?
* `embed_size` dimensionality of word vectors (use 300 for our pretrained vectors)
* `hidden_size` dimensionality of hidden layers (e.g. in the encoder)
* `num_layers` number of encoder layers (e.g. you may want to stack LSTMs)
* `layer_inf` which encoder should be used in the inference network (e.g. bow/lstm)
* `layer_cls` which composition function should the classifier use (e.g. bow/lstm)

```python
import torch.optim
# and a couple of tricks to reduce learning rate on plateau
from torch.optim.lr_scheduler import ReduceLROnPlateau

cfg = dict()

# Data (you probably don't need to change these)
cfg['training_path'] = "data/sst/train.txt"
cfg['dev_path'] = "data/sst/dev.txt"
cfg['test_path'] = "data/sst/test.txt"
cfg['word_vectors'] = 'data/sst/glove.840B.300d.filtered.txt'
cfg['subphrases'] = False
cfg['min_phrase_length'] = 2
cfg['lowercase'] = True
# Model
cfg['prior_p1'] = 0.3
# Architecture
cfg['fix_emb'] = True
cfg['embed_size'] = 300
cfg['hidden_size'] = 150
cfg['num_layers'] = 1
cfg['dropout'] = 0.5
cfg['layer_inf'] = 'bow'
cfg['layer_cls'] = 'bow'
# Posterior collapse
cfg['min_kl'] = 2.  # use more than 0 to enable free bits
cfg['kl_weight'] = 0.001  # start from zero to enable annealing
cfg['kl_inc'] = 0.0001
# Logging
cfg['save_path'] = 'data/results'
cfg['print_every'] = 100
cfg['eval_every'] = -1
cfg['batch_size'] = 25
cfg['eval_batch_size'] = 25
cfg['display_dev'] = [39, 41, 1, 0, 4]  # range(10) or selected examples
# Optimiser (leave as is)
cfg['num_epochs'] = 100
cfg['running_alpha'] = 0.9
cfg['gen_opt'] = 'adam'
cfg['gen_lr'] = 0.001
cfg['gen_momentum'] = 0.
cfg['inf_opt'] = 'adadelta'
cfg['inf_lr'] = 0.0001
cfg['inf_momentum'] = 0.
cfg['gen_weight_decay'] = 1e-5
cfg['inf_weight_decay'] = 1e-5
cfg['lr_decay'] = 0.5
cfg['patience'] = 10
cfg['cooldown'] = 10
cfg['threshold'] = 1e-4
cfg['min_lr'] = 1e-6
cfg['max_grad_norm'] = 5.

print('# Configuration')
for k, v in cfg.items():
    print("{:20} : {:10}".format(k, str(v)))

iters_per_epoch = len(train_data) // cfg["batch_size"]

if cfg["eval_every"] == -1:
    eval_every = iters_per_epoch
    print("Set eval_every to {}".format(iters_per_epoch))
```

Here we load the data and display an example

```python
# Let's load the data into memory.
print("Loading data")
train_data = list(examplereader(
    cfg['training_path'],
    lower=cfg['lowercase'],
    subphrases=cfg['subphrases'],
    min_length=cfg['min_phrase_length']))
dev_data = list(examplereader(cfg['dev_path'], lower=cfg['lowercase']))
test_data = list(examplereader(cfg['test_path'], lower=cfg['lowercase']))

print("train", len(train_data))
print("dev", len(dev_data))
print("test", len(test_data))

print('\n# Example')
example = dev_data[0]
print("First dev example:", example)
print("First dev example tokens:", example.tokens)
print("First dev example label:", example.label)
```

Here we realise the main training loop.

```python
from functools import partial
import torch.optim as optim

def update_avg(old, new, alpha):
    """Update a running average"""
    if old is None:
        return new
    return old * alpha + (1 - alpha) * new

def get_optimizer(name, parameters, lr, l2_weight, momentum=0.):
    if name is None or name == "adam":
        cls = optim.Adam
    elif name == "amsgrad":
        cls = partial(optim.Adam, amsgrad=True)
    elif name == "adadelta":
        cls = optim.Adadelta
    elif name == "rmsprop":
        cls = partial(optim.RMSprop, momentum=momentum)
    elif name == 'sgd':
        cls = optim.SGD
    else:
        raise ValueError("Unknown optimizer: %s" % name)
    return cls(params=parameters, lr=lr, weight_decay=l2_weight)

def train():

    # Create a vocabulary object to map str <-> int
    vocab = Vocabulary()
    # populate it with pretrained word vectors
    glove_path = cfg["word_vectors"]
    vectors = load_glove(glove_path, vocab)

    # You may consider using tensorboardX
    # writer = SummaryWriter(log_dir=cfg["save_path"])

    # Map the sentiment labels 0-4 to a more readable form (and the opposite)
    i2t = ["very negative", "negative", "neutral", "positive", "very positive"]
    t2i = OrderedDict({p: i for p, i in zip(i2t, range(len(i2t)))})

    print('\n# Constructing model')
    model = VAE(
        vocab_size=len(vocab.w2i),
        emb_size=cfg["embed_size"],
        hidden_size=cfg["hidden_size"],
        num_classes=len(t2i),
        prior_p1=cfg['prior_p1'],
        vocab=vocab,
        dropout=cfg["dropout"],
        layer_cls=cfg["layer_cls"],
        layer_inf=cfg["layer_inf"])

    print('\n# Loading embeddings')
    with torch.no_grad():
        model.embed.weight.data.copy_(torch.from_numpy(vectors))
        if cfg["fix_emb"]:
            print("fixed word embeddings")
            model.embed.weight.requires_grad = False
        model.embed.weight[1] = 0.  # padding zero

    # Configure optimizers
    #  it's not uncommon to use different optimisers when we employ discrete latent variables
    #  the reason is because the score function estimator (SFE) and the reparameterised gradient
    #  have very difference variances, and it's useful to have smaller learning rates for the SFE
    gen_optimizer = get_optimizer(
        cfg['gen_opt'],
        model.generative_parameters(),
        cfg['gen_lr'], cfg['gen_weight_decay'], cfg['gen_momentum']
    )
    inf_optimizer = get_optimizer(
        cfg['inf_opt'],
        model.inference_parameters(),
        cfg['inf_lr'], cfg['gen_weight_decay'], cfg['inf_momentum']
    )

    # and learning rate scheduler
    gen_scheduler = ReduceLROnPlateau(
        gen_optimizer, mode="min", factor=cfg["lr_decay"], patience=cfg["patience"],
        verbose=True, cooldown=cfg["cooldown"], threshold=cfg["threshold"],
        min_lr=cfg["min_lr"])
    inf_scheduler = ReduceLROnPlateau(
        inf_optimizer, mode="min", factor=cfg["lr_decay"], patience=cfg["patience"],
        verbose=True, cooldown=cfg["cooldown"], threshold=cfg["threshold"],
        min_lr=cfg["min_lr"])

    # Prepare a few auxiliary variables
    iter_i = 0
    best_eval = 1.0e9
    best_iter = 0

    model = model.to(device)

    # Some debugging info
    print(model)
    print_parameters(model)

    batch_size = cfg['batch_size']
    eval_batch_size = cfg['eval_batch_size']
    print_every = cfg['print_every']

    # Parameters of tricks to better optimise the ELBO
    kl_inc = cfg['kl_inc']
    kl_weight = cfg['kl_weight']
    min_kl = cfg['min_kl']
    # Running estimates for baselines
    running_avg, running_std = 0., 1.
    running_elbo, running_kl, running_kl_ = None, None, None
    running_loss, running_reward, running_rate = None, None, None
    alpha = cfg['running_alpha']

    while True:  # when we run out of examples, shuffle and continue
        for batch in get_minibatch(train_data, batch_size=batch_size, shuffle=True):

            epoch = iter_i // iters_per_epoch
            if epoch > cfg['num_epochs']:
                break

            # forward pass
            model.train()
            x, y, _ = prepare_minibatch(batch, model.vocab, device=device)
            mask = (x != 1)

            # "KL annealing"
            kl_weight += kl_inc
            if kl_weight > 1.:
                kl_weight = 1.0

            loss, terms = model.get_loss(
                x, mask, y,
                kl_weight=kl_weight,
                min_kl=min_kl,
                ll_mean=running_avg,
                ll_std=running_std,
                iter_i=iter_i)

            # backward pass
            model.zero_grad()  # erase previous gradients
            loss.backward()  # compute new gradients
            # gradient clipping generally helps
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=cfg['max_grad_norm'])
            # update weights
            gen_optimizer.step()
            inf_optimizer.step()

            # logging
            running_loss = update_avg(running_loss, loss.item(), alpha)
            running_elbo = update_avg(running_elbo, terms['elbo'], alpha)
            running_kl = update_avg(running_kl, terms['kl'], alpha)
            running_kl_ = update_avg(running_kl_, terms['kl_'], alpha)
            running_reward = update_avg(running_reward, terms['reward'], alpha)
            running_rate = update_avg(running_rate, terms['selected'], alpha)

            # keep an running estimate of the reward (log P(y|x,z))
            running_avg = update_avg(running_avg, terms['ll'], alpha)
            running_std = update_avg(running_std, terms['ll_std'], alpha)
            if running_std < 1.:
                running_std = 1.

            iter_i += 1

            # print info
            if iter_i % print_every == 0:
                print("Epoch %r Iter %r loss=%.4f ELBO (KL/KL_)=%.4f (%.2f/%.2f) kl_weight=%.4f reward=%.4f selected=%.2f" %
                      (epoch, iter_i, running_loss,
                       running_elbo, running_kl, running_kl_,
                       kl_weight, running_reward,
                       running_rate
                      ))

            # evaluate
            if iter_i % eval_every == 0:

                dev_elbo, dev_kl = eval_elbo(
                    model, dev_data,
                    batch_size=eval_batch_size,
                    device=device,
                    iter_i=0,
                    nb_samples=10)
                acc, selection_rate, rationales = evaluate(
                    model, dev_data,
                    batch_size=eval_batch_size,
                    device=device,
                    iter_i=iter_i)

                print("\n# epoch %r iter %r: dev ELBO (KL) %.4f (%.2f) acc %.2f selected %.2f" % (
                    epoch, iter_i, dev_elbo, dev_kl, acc, selection_rate ))

                for exid in cfg['display_dev']:
                    print(' dev%d [gold=%d,pred=%d]:' % (exid, dev_data[exid].label, rationales[exid][1]),
                          ' '.join(rationales[exid][0]))
                print()

                # adjust learning rate
                gen_scheduler.step(-dev_elbo)
                inf_scheduler.step(-dev_elbo)
```

```python
train()
```

# Variance reduction

**This is an extra**

We can use a *control variate* to reduce the variance of our gradient estimates.

Let's recap the idea in general terms. We are looking to solve some expectation
\begin{align}
\mu_f = \mathbb E[f(Z)]
\end{align}
but unfortunatelly, realising the full sum (or integral for continuous variables) is intractable. Thus we employ MC estimation
\begin{align}
\hat \mu_f &\overset{\text{MC}}{\approx} \frac{1}{S} \sum_{s=1}^S f(z_s) & \text{where }z_s \sim Q(z|x)
\end{align}
Note that the variance of this estimate is
\begin{align}
\text{Var}(\hat \mu_f) &=  \frac{1}{S}\text{Var}(f(Z)) \\
&= \frac{1}{S} \mathbb E[( f(Z) - \mathbb E[f(Z)])^2]
\end{align}
Note that this variance is such that it goes down as we sample more, in a rate $\mathcal O(S^{-1})$.
See that if we sample $10$ times more, we will only obtain an decrease in variance in the order of $10^{-1}$. This means that sampling more is generally not the most convenient way to decrease variance.

*Digression* we can estimate the variance itself via MC, an unbiased estimate looks like
\begin{align}
\hat \sigma^2_f = \frac{1}{S(S-1)} \sum_{s=1}^S (f(z_s) - \hat \mu_f)^2
\end{align}
but not that this estimate is even hard to improve since it decreases with $\mathcal O(S^{-2})$.

Back to our main problem: let's try and improve the variance of our estimator to $\mu_f$.

It's a fact, and it can be shown trivially, that
\begin{align}
\mu_f &=  \mathbb E[f(Z) - \psi(Z)] + \underbrace{\mathbb E[\psi(Z)]}_{\mu_\psi} \\
 &\overset{\text{MC}}{\approx} \underbrace{\left(\frac{1}{S} \sum_{s=1}^S f(z_s) - \psi(z_s) \right) + \mu_\psi}_{\hat c}
\end{align}
where we assume the existence of some function $\psi(z)$ for which the expected value $\mu_\psi$ is known and we estimate the expected difference $\mathbb E[f(Z) - \psi(Z)]$ via MC. We used this auxiliary function, also known as a *control variate*, to derive a new estimator, which we will denote by $\hat c$.

The variance of this new estimator is shown below:

\begin{align}
\text{Var}( \hat c ) &= \text{Var}(\hat \mu_{f-\psi}) + 2\underbrace{\text{Cov}(\hat \mu_{f-\psi}, \mu_\psi)}_{\mathbb E[\hat \mu_{f-\psi}  \mu_\psi] - \mathbb E[\hat \mu_{f-\psi}] \mathbb E[\mu_\psi]} + \underbrace{\text{Var}(\mu_\psi)}_{\color{blue}{0} } \\
&= \frac{1}{S}\text{Var}(f- \psi)  + 2 \underbrace{\left( \mu_\psi \mu_{f-\psi} - \mu_{f-\psi} \mu_\psi \right)}_{\color{blue}{0}}
\end{align}
where the variance of $\mu_\psi$ is 0 because we know it in closed form (no need for MC estimation), and the covariance is $0$ as shown in the second row.

That is, the variance of $\hat c$ is essentially the variance of estimating $\mathbb E[f(Z) - \psi(Z)]$, which in turn depends on the variance

\begin{align}
\text{Var}(f-\psi) &= \text{Var}(f) - 2\text{Cov}(f, \psi) + \text{Var}(\psi)
\end{align}
where we can see that if $\text{Cov}(f, \psi) > \frac{\text{Var}(\psi)}{2}$ we achieve variance reduction as then $\text{Var}(f-\psi)$ would be smaller than $\text{Var(f)}$.

## Baselines

Baslines are control variates of a very simple form:
\begin{align}
\mathbb E[f(Z)] = \mathbb E[f(Z) - C] + \mathbb E[C]
\end{align}
where $C$ is a constant with respect to $z$.

In the context of the score function estimator, a baseline looks like a quantity $C(x; \omega)$, this may be
* just a constant;
* or a function of the input (but not of the latent variable), which could be itself implemented as a neural network;
* a combination of the two.

Let's focus on the first term of the ELBO (so I'm omitting the KL term here). The gradient with respect to parameters of the inference model becomes:

\begin{align}
&\mathbb E_{Q(z|x, \lambda)}\left[ \log P(x|z, \theta) \nabla_\lambda \log Q(z|x, \lambda)\right]\\
&=\mathbb E_{Q(z|x, \lambda)}\left[\log P(x|z, \theta) \nabla_\lambda \log Q(z|x, \lambda) - \color{red}{C(x; \omega)\nabla_\lambda \log Q(z|x, \lambda) }  \right] + \underbrace{\mathbb E_{Q(z|x, \lambda)}\left[\color{red}{C(x; \omega)\nabla_\lambda \log Q(z|x, \lambda) }  \right] }_{=0} \\
&= \mathbb E_{Q(z|x, \lambda)}\left[ \color{blue}{\left(\log P(x|z, \theta) - C(x; \omega) \right)}\nabla_\lambda \log Q(z|x, \lambda)\right] \\
&
\end{align}
We can show that the last term is $0$

\begin{align}
&\mathbb E_{Q(z|x, \lambda)}\left[C(x; \omega)\nabla_\lambda \log Q(z|x, \lambda)   \right]  \\&= C(x; \omega) \mathbb E_{Q(z|x, \lambda)}\left[\nabla_\lambda \log Q(z|x, \lambda)   \right]\\
&= C(x; \omega) \mathbb E_{Q(z|x, \lambda)}\left[\frac{1}{Q(z|x, \lambda)} \nabla_\lambda Q(z|x, \lambda)   \right] \\
&= C(x; \omega) \sum_z Q(z|x, \lambda) \frac{1}{Q(z|x, \lambda)} \nabla_\lambda Q(z|x, \lambda)   \\
&= C(x; \omega) \sum_z\nabla_\lambda Q(z|x, \lambda)  \\
&= C(x; \omega) \nabla_\lambda \underbrace{\sum_z Q(z|x, \lambda)  }_{=1}\\
&=0
\end{align}

Examples of useful baselines:

* a running average of the learning signal: at some iteration $t$ we can use a running average of $\log P(x|z, \theta)$ using parameter estimates $\theta$ from iterations $i < t$, this is a baseline that likely leads to high correlation between control variate and learning signal and can lead to variance reduction;
* another technique is to have an MLP with parameters $\omega$ predict a scalar and train this MLP to approximate the learning signal $\log P(x|z, \theta)$ via regression:
\begin{align}
\arg\max_\omega \left( C(x; \omega) - \log P(x|z, \theta) \right)^2
\end{align}
its left as an extra to implement these ideas.

One more note: we can also use something called a *multiplicative baseline* in the literature of reinforcement learning, whereby we incorporate a running estimate of the standard deviation of the learning signal computed based on the values attained on previous iterations:
\begin{align}
\mathbb E_{Q(z|x, \lambda)}\left[ \frac{1}{\hat\sigma_{\text{past}}}\left(\log P(x|z, \theta) - \hat \mu_{\text{past}}\right)\nabla_\lambda \log Q(z|x, \lambda)\right]
\end{align}
this form of contorl variate aim at promoting the learning signal (or reward in reinforcement learning literature) to be distributed by $\mathcal N(0, 1)$. Note that multiplying the reward by a constant does not bias the estimator, and in this case, may lead to variance reduction.

```python

```

```python

```


# Mooney-Rivlin Hyperelasticity

## Overview

A Mooney-Rivlin hyperelastic material is one for which the derivatives of the free energy with respect to the invariants of stretch are constant.  The Mooney-Rivlin model

- a special case of the more general [polynomial hyperelastic model](PolynomialHyperelastic.ipynb)
- has 2 elastic constants, typically fit to uniaxial extension/compression, equibiaxial extension/compression, or shear experimental data
- usually valid for strains less than 100%

## See Also

- [User Defined Materials](UserMaterial.ipynb)
- [Linear Elastic Material](LinearElastic.ipynb)
- [Polynomial Hyperelastic Material](PolynomialHyperelastic.ipynb)

## Contents

1. Fundamental Equations
2. Model Implementation
3. Model Verification

```python
from bokeh.io import output_notebook
from bokeh.plotting import *
from matmodlab2 import *
from numpy import *
import numpy as np
from plotting_helpers import create_figure
output_notebook()
```

    Setting up the Matmodlab notebook environment

    Loading BokehJS ...

## Fundamental Equations

The Mooney-Rivlin material is a special case of the more general polynomial hyperelastic model defined by the following free energy potential

$$
a = c_{10}\left(\overline{I}_1 - 3\right) +  c_{01}\left(\overline{I}_2 - 3\right)
  + \frac{1}{D_1}\left(J-1\right)^2
$$

where $c_{10}$, $c_{01}$, and $D_1$ are material parameters, the $\overline{I}_i$ are the isochoric invariants of the right Cauchy deformation tensor $C_{IJ} = F_{kI}F_{kJ}$, $F_{iJ}$ are the components of the deformation gradient tensor, and $J$ is the determinant of the deformation gradient.

### Second Piola-Kirchhoff Stress

The components of the second Piola-Kirchhoff stress $S_{IJ}$ are given by

$$
\frac{1}{2\rho_0} S_{IJ} = \frac{\partial a}{\partial C_{IJ}}
$$

For an isotropic material, the free energy $a$ is a function of $C_{IJ}$ through only its invariants, giving for $S_{IJ}$

$$
\frac{1}{2\rho_0}S_{IJ}
       = \frac{\partial a}{\partial\overline{I}_1}\frac{\partial\overline{I}_1}{\partial C_{IJ}}
       + \frac{\partial a}{\partial\overline{I}_2}\frac{\partial\overline{I}_2}{\partial C_{IJ}}
       + \frac{\partial a}{\partial J}\frac{\partial J}{\partial C_{IJ}}
$$

The partial derivatives of the energy with respect to the invariants are

$$
\frac{\partial a}{\partial\overline{I}_1} = c_{10}, \quad
\frac{\partial a}{\partial\overline{I}_2} = c_{01}, \quad
\frac{\partial a}{\partial J} = \frac{2}{D_1}\left(J-1\right)
$$

and the partials of the invariants with respect to $C_{IJ}$ are

$$
\begin{align}
\frac{\partial\overline{I}_1}{\partial C_{IJ}}
 &= \frac{\partial\left( J^{-2/3} C_{KK}\right)}{\partial C_{IJ}} \\
 %&= \frac{\partial J^{-2/3}}{\partial C_{IJ}} I_1 + J^{-2/3} \frac{\partial C_{KK}}{\partial C_{IJ}} \\
 %&= -\frac{2}{3}J^{-5/3}\frac{\partial J}{\partial C_{IJ}} I_1 + J^{-2/3} \delta_{IJ} \\
 %&= -\frac{2}{3}J^{-5/3}\frac{1}{2}J C_{IJ}^{-1} I_1 + J^{-2/3} \delta_{IJ} \\
 &= J^{-2/3}\left(\delta_{IJ} - \frac{1}{3}I_1 C_{IJ}^{-1}\right) \\
\end{align}
$$

$$
\begin{align}
\frac{\partial\overline{I}_2}{\partial C_{IJ}}
 &= \frac{\partial\left( J^{-4/3}\left(C_{KK}^2 - C_{MN}C_{MN}\right)\right)}{\partial C_{IJ}} \\
 %&= \frac{\partial J^{-4/3}}{\partial C_{IJ}} \left(C_{KK}^2 - C_{MN}C_{MN}\right) + J^{-4/3} \frac{\partial \left(C_{KK}^2 - C_{MN}C_{MN}\right)}{\partial C_{IJ}} \\
 &= J^{-4/3}\left(I_1\delta_{IJ} - C_{IJ} - \frac{2}{3} I_2 C_{IJ}^{-1}\right) \\
\end{align}
$$

$$
\begin{align}
\frac{\partial J}{\partial C_{IJ}} &= \frac{\partial \sqrt{\det C_{KL}}}{\partial C_{IJ}} \\
%&= \frac{1}{2 \sqrt{\det C_{KL}}} \frac{\partial\det C_{KL}}{\partial C_{IJ}} \\
%&= \frac{1}{2 J} \frac{\partial\det C_{KL}}{\partial C_{IJ}} \\
&= \frac{1}{2} J C_{IJ}^{-1} \\
\end{align}
$$

Combining the above results, we arrive at the following expression for the components $S_{IJ}$:

$$
\frac{1}{2\rho_0}S_{IJ}
   = c_{10}J^{-2/3}\left(\delta_{IJ} - \frac{1}{3}I_1 C_{IJ}^{-1}\right)
   + c_{01}J^{-4/3}\left(I_1\delta_{IJ} - C_{IJ} - \frac{2}{3} I_2 C_{IJ}^{-1}\right)
   + \frac{J}{D_1}\left(J-1\right) C_{IJ}^{-1}
$$

### Cauchy Stress

The components of the Cauchy stress tensor $\sigma_{ij}$ are given by the push-forward $S_{IJ}$:

$$
J \sigma_{ij} = F_{iM}S_{MN}F_{jN}
$$

$$
\begin{align}
\frac{J}{2\rho_0}\sigma_{ij} &= F_{iM}\left[
     c_{10}J^{-2/3}\left(\delta_{MN} - \frac{1}{3}I_1 C_{MN}^{-1}\right)
   + c_{01}J^{-4/3}\left(I_1\delta_{MN} - C_{MN} - \frac{2}{3} I_2 C_{MN}^{-1}\right)
   + \frac{J}{D_1}\left(J-1\right) C_{MN}^{-1}\right] F_{jN} \\
  &= c_{10}J^{-2/3}\left(F_{iN}F_{jN} - \frac{1}{3}I_1 F_{iM}C_{MN}^{-1}F_{jN}\right)
   + c_{01}J^{-4/3}\left(I_1F_{iN}F_{jN} - F_{iM}C_{MN}F_{jN} - \frac{2}{3} I_2 F_{iM}C_{MN}^{-1}F_{jN}\right) \\
   &\qquad + \frac{J}{D_1}\left(J-1\right) F_{iM}C_{MN}^{-1}F_{jN} \\
\end{align}
$$

Recognizing that

$$
C_{MN} = F_{kM}F_{kN} \Rightarrow C_{MN}^{-1} = F_{kN}^{-1}F_{kM}^{-1}
$$

the components of the Cauchy stress can be written as

$$
\begin{align}
\frac{J}{2\rho_0}\sigma_{ij}
  &= c_{10}J^{-2/3}\left(F_{iN}F_{jN} - \frac{1}{3}I_1 F_{iM}F_{kN}^{-1}F_{kM}^{-1}F_{jN}\right)
   + c_{01}J^{-4/3}\left(I_1F_{iN}F_{jN} - F_{iM}F_{kM}F_{kN}F_{jN} - \frac{2}{3} I_2 F_{iM}F_{kN}^{-1}F_{kM}^{-1}F_{jN}\right) \\
   &\qquad + \frac{J}{D_1}\left(J-1\right) F_{iM}F_{kN}^{-1}F_{kM}^{-1}F_{jN} \\
  &= c_{10}J^{-2/3}\left(B_{ij} - \frac{1}{3}I_1 \delta_{ij} \right)
   + c_{01}J^{-4/3}\left(I_1B_{ij} - B_{ik}B_{kj} - \frac{2}{3} I_2 \delta_{ij}\right)
   + \frac{J}{D_1}\left(J-1\right) \delta_{ij} \\
  &= J^{-2/3} \left(c_{10} + c_{01}\overline{I}_1 \right)B_{ij}
   - c_{01}J^{-4/3}B_{ik}B_{kj}
  + \left(
      \frac{J}{D_1}\left(J-1\right)  - \frac{1}{3}\left(c_{10}\overline{I}_1 + c_{01}2\overline{I}_2\right)
    \right)\delta_{ij}
\end{align}
$$

where $B_{ij} = F_{iN}F_{jN}$ is the left Cauchy deformation tensor.  Finally, the components of the Cauchy stress are given by

$$
\begin{align}
\sigma_{ij} &= \frac{2\rho_0}{J} \left(
  J^{-2/3} \left(c_{10} + c_{01}\overline{I}_1 \right)B_{ij}
   - c_{01}J^{-4/3}B_{ik}B_{kj}\right)
  + \left(
      \frac{2\rho_0}{D_1}\left(J-1\right)  - \frac{2\rho_0}{3J}\left(c_{10}\overline{I}_1 + c_{01}2\overline{I}_2\right)
    \right)\delta_{ij} \\
&= \frac{2\rho_0}{J} \left(
  \left(c_{10} + c_{01}\overline{I}_1 \right)\overline{B}_{ij}
   - c_{01}\overline{B}_{ik}\overline{B}_{kj}\right)
  + \left(
      \frac{2\rho_0}{D_1}\left(J-1\right)  - \frac{2\rho_0}{3J}\left(c_{10}\overline{I}_1 + c_{01}2\overline{I}_2\right)
    \right)\delta_{ij} \\
\end{align}
$$

where $\overline{B}_{ij} = J^{-2/3}B_{ij}$

## Matmodlab Implementation

The Mooney-Rivlin material is implemented as the `MooneyRivlinMaterial` class in `matmodlab2/materials/mooney_rivlin.py`.  The file can be viewed by executing the following cell.

```python
%pycat ../matmodlab2/materials/mooney_rivlin.py
```

## Verification

### Uniaxial Stress

For an incompressible isotropic material, uniaxial stress is produced by the following deformation state

$$
[F] = \begin{bmatrix}\lambda && \\ & \frac{1}{\sqrt{\lambda}} & \\ & & \frac{1}{\sqrt{\lambda}} \end{bmatrix}
$$

The stress difference $\sigma_{\text{axial}} - \sigma_{\text{lateral}}$ is given by

```python
from sympy import Symbol, Matrix, Rational, symbols, sqrt
lam = Symbol('lambda')
F = Matrix(3, 3, [lam, 0, 0, 0, 1/sqrt(lam), 0, 0, 0, 1/sqrt(lam)])
B = Matrix(3, 3, F.dot(F.T))
Bsq = Matrix(3, 3, B.dot(B))
I = Matrix(3, 3, lambda i,j: 1 if i==j else 0)
I1 = B.trace()
I2 = ((B.trace()) ** 2 - Bsq.trace()) / 2
J = F.det()
X = J ** Rational(1, 3)
C1, C2, D1 = symbols('C10 C01 D1')
I1B = I1 / X ** 2
I2B = I2 / X ** 4

S = 2 / J * (1 / X ** 2 * (C1 + I1B * C2) * B - 1 / X ** 4 * C2 * Bsq) \
    + (2 / D1 * (J - 1) - 2 * (C1 * I1B + 2 * C2 * I2B) / 3) * I
(S[0,0] - S[1,1]).simplify()
```

    2*C01*lambda - 2*C01/lambda**2 + 2*C10*lambda**2 - 2*C10/lambda

We now exercise the Mooney-Rivlin material model using Matmodlab

```python
# Hyperelastic parameters, D1 set to a large number to force incompressibility
parameters = {'D1': 1.e12, 'C10': 1e6, 'C01': .1e6}

# stretch to 300%
lam = linspace(.5, 3, 50)

# Set up the simulator
mps = MaterialPointSimulator('test1')
mps.material = MooneyRivlinMaterial(**parameters)

# Drive the *incompressible* material through a path of uniaxial stress by
# prescribing the deformation gradient.
Fij = lambda x: (x, 0, 0, 0, 1/sqrt(x), 0, 0, 0, 1/sqrt(x))
mps.run_step('F', Fij(lam[0]), frames=10)
mps.run_step('F', Fij(1), frames=1)
mps.run_step('F', Fij(lam[-1]), frames=20)

# plot the analytic solution and the simulation
p = create_figure(bokeh=True, x_axis_label='Stretch', y_axis_label='Stress')
C10, C01 = parameters['C10'], parameters['C01']

# analytic solution for true and engineering stress
s = 2*C01*lam - 2*C01/lam**2 + 2*C10*lam**2 - 2*C10/lam

# plot the analytic solutions
p.line(lam, s, color='blue', legend='True', line_width=2)
p.line(lam, s/lam, color='green', legend='Engineering', line_width=2)

lam_ = np.exp(mps.get('E.XX'))
ss = mps.get('S.XX') - mps.get('S.ZZ')
p.circle(lam_, ss, color='orange', legend='Simulation, True')
p.circle(lam_, ss/lam_, color='red', legend='Simulation, Engineering')
p.legend.location = 'top_left'

show(p)

# check the actual solutions
assert abs(amax(ss) - amax(s)) / amax(s) < 1e-6
assert abs(amin(ss) - amin(s)) < 1e-6
```

```python

```


This notebook is sample of the HAL QCD potential,
the effective mass fitting, and the effective energy shifts of two-baryon system
from compressed NBS wavefunction sample_data.

In order to decompress the wave function, hal_pot_single_ch.py requires
binary "PH1.compress48" in "yukawa library."

```python
%pylab inline
```

    Populating the interactive namespace from numpy and matplotlib

```python
import seaborn as sns
sns.set_style('ticks', {'axes.grid': True})
sns.set_context('poster', font_scale=2.0)
%config InlineBackend.figure_format = 'retina'

plt.rcParams['figure.figsize'] = (12.8, 9.6)
plt.rcParams['figure.facecolor'] = 'white'

```

```python
ls ../data/sample_data/
```

    BBwave.dir.S4.00/     correlator.PS.dir/
    README.md             correlator.multi.dir/

```python
# import library
from hal_pot_single_ch import HAL_pot
from corr_baryon import Corr_2pt_Baryon, Corr_2pt_2Baryons, Delta_Eeff
```

lattice spacing

```python
ainv = 2.194e3 # MeV
hbarc = 0.197e3 # MeV fm
lat_unit = hbarc/ainv
```

analyze baryon correlator and mass
\begin{equation}
m_\mathrm{eff}(t) = \log \frac{C_\mathrm{B}(t)}{C_\mathrm{B}(t+1)}
\end{equation}

```python
cb = Corr_2pt_Baryon('Xi_CG05_CG05', bin_size=1, result_dir='../data/sample_data/')
```

    # Read Xi_CG05_CG05
    # total 4 conf., bin size = 1, number of samples =  4

```python
cb.plot_meff()
```

```python
fig, ax = plt.subplots()
cb.fit_meff(fit_min=15, fit_max=20, ax=ax)
ax.set_ylim(0.65, 0.68)
leg = ax.legend(frameon=True)
leg.get_frame().set_edgecolor('black')
leg.get_frame().set_linewidth(2.0)
ax.set_title('$\Xi$-baryon effective mass');
```

```python
m_red = 0.5 * 0.665
```

initialize HAL_pot object with induced mass as an input parameter

```python
hal = HAL_pot(m_red=m_red, result_dir='../data/sample_data/',
              it0 = 10, Ns=48, channel='xixi', bin_size=2)
```

    # Read Xi_CG05_CG05
    # total 4 conf., bin size = 2, number of samples =  2

calc_pot method evaluate (effective) central potential of (3S1) 1S0 channel, and central, tensor part of 3S1 channel.

```python
# potential
result = hal.calc_pot(pot_type='cen', spin='1s0')
```

    # calc. Rcorr  results.rcorr.binned/Rcorr_1s0_xixi_t009_2bin_4conf.dat
    # binning and decompress NBS
    # binning NBS S4.00 t = 9
    --- binning NBS wavefunc. --- 000
    >> decompress results.nbs.binned/NBSwave.S4.00.t009.binned.000.dat
    --- binning NBS wavefunc. --- 001
    >> decompress results.nbs.binned/NBSwave.S4.00.t009.binned.001.dat
    >> total 2 binned NBS files
    # load  results.nbs.decomp/NBSwave.S4.00.t009.binned.000.decomp.dat
    # load  results.nbs.decomp/NBSwave.S4.00.t009.binned.001.decomp.dat
    # save results.rcorr.binned/Rcorr_1s0_xixi_t009_2bin_4conf.dat
    # calc. Rcorr  results.rcorr.binned/Rcorr_1s0_xixi_t010_2bin_4conf.dat
    # binning and decompress NBS
    # binning NBS S4.00 t = 10
    --- binning NBS wavefunc. --- 000
    >> decompress results.nbs.binned/NBSwave.S4.00.t010.binned.000.dat
    --- binning NBS wavefunc. --- 001
    >> decompress results.nbs.binned/NBSwave.S4.00.t010.binned.001.dat
    >> total 2 binned NBS files
    # load  results.nbs.decomp/NBSwave.S4.00.t010.binned.000.decomp.dat
    # load  results.nbs.decomp/NBSwave.S4.00.t010.binned.001.decomp.dat
    # save results.rcorr.binned/Rcorr_1s0_xixi_t010_2bin_4conf.dat
    # calc. Rcorr  results.rcorr.binned/Rcorr_1s0_xixi_t011_2bin_4conf.dat
    # binning and decompress NBS
    # binning NBS S4.00 t = 11
    --- binning NBS wavefunc. --- 000
    >> decompress results.nbs.binned/NBSwave.S4.00.t011.binned.000.dat
    --- binning NBS wavefunc. --- 001
    >> decompress results.nbs.binned/NBSwave.S4.00.t011.binned.001.dat
    >> total 2 binned NBS files
    # load  results.nbs.decomp/NBSwave.S4.00.t011.binned.000.decomp.dat
    # load  results.nbs.decomp/NBSwave.S4.00.t011.binned.001.decomp.dat
    # save results.rcorr.binned/Rcorr_1s0_xixi_t011_2bin_4conf.dat
    >> return jackknife samples of (effective) central
    # output potential >>  results.pot/pot_1s0_cen_xixi_t010_004conf_002bin.dat

text data of potential is stored in results.pot dir.

```python
pot = np.loadtxt('results.pot/pot_1s0_cen_xixi_t010_004conf_002bin.dat')
```

```python
pot.shape
```

    (2925, 9)

columns of potential data

0, 1, 2, 3, 4, 5, 6, 7, 8

r, H0-term (av, err) , d/dt-term (av, err), d2/dt2-term (av, err), total (av, err)

```python
fig, ax = plt.subplots()
ax.errorbar(pot[:,0]*lat_unit, pot[:,1]*ainv, pot[:,2]*ainv,
            fmt='bs', mfc='none', mew=2, capsize=10, capthick=2, label=r'$H_0$-term')

ax.errorbar(pot[:,0]*lat_unit, pot[:,3]*ainv, pot[:,4]*ainv,
           fmt='g^', mfc='none', mew=2, capsize=10, capthick=2, label=r'$\partial/\partial t$-term')

ax.errorbar(pot[:,0]*lat_unit, pot[:,5]*ainv, pot[:,6]*ainv,
           fmt='kd', mfc='none', mew=2, capsize=10, capthick=2, label=r'$\partial^2/\partial t^2$-term')

ax.errorbar(pot[:,0]*lat_unit, pot[:,7]*ainv, pot[:,8]*ainv,
           fmt='ro', mfc='none', mew=2, capsize=10, capthick=2, label='total')

ax.set_ylim(-50, 50)

ax.axhline(0, color='black')
ax.set_xlabel(r'$r$ [fm]', size=48)
ax.set_ylabel(r'$V(r)$ [MeV]', size=48)
leg = ax.legend(frameon=True)
leg.get_frame().set_edgecolor('black')
leg.get_frame().set_linewidth(2.0)
```

# baryon mass and energy shift

if lattice unit is given, observables are plotted in physical scale

```python
cb = Corr_2pt_Baryon('Xi_CG05_CG05', bin_size=1, result_dir='../data/sample_data/', lat_unit=lat_unit)
```

    # Read Xi_CG05_CG05
    # total 4 conf., bin size = 1, number of samples =  4

```python
fig, ax = plt.subplots()
cb.fit_meff(fit_min=10, fit_max=20, ax=ax)
leg = ax.legend(frameon=True)
leg.get_frame().set_edgecolor('black')
leg.get_frame().set_linewidth(2.0)
ax.set_ylim(1400, 1500)
```

Two-baryon correlator and the effective energy
\begin{equation}
E_\mathrm{eff}(t) = \log \frac{C_\mathrm{BB}(t)}{C_\mathrm{BB}(t+1)}
\end{equation}

```python
# BB two baryon correlator
cbb = Corr_2pt_2Baryons('xixi', bin_size=1, result_dir='../data/sample_data/', lat_unit=lat_unit)
```

    Read xixi
    # total 4 conf., bin size = 1, number of samples =  4

```python
fig, ax = plt.subplots()
cbb.fit_Eeff(ax=ax)
ax.set_ylim(2800, 3000)
```

effective energy shift
\begin{equation}
\Delta E_\mathrm{eff}(t) = \log \frac{R(t)}{R(t+1)}
\end{equation}
with
\begin{equation}
R(t) = \frac{C_\mathrm{BB}(t)}{\{C_\mathrm{B}(t)\}^2}
\end{equation}

```python
# effective energy shifts
dEeff = Delta_Eeff('xixi',  bin_size=1, result_dir='../data/sample_data', lat_unit=lat_unit)
```

    Read xixi
    # total 4 conf., bin size = 1, number of samples =  4
    # Read Xi_CG05_CG05
    # total 4 conf., bin size = 1, number of samples =  4

```python
fig, ax = plt.subplots()
dEeff.fit_dEeff(fit_min=7, fit_max=10, spin='1s0', ax=ax)
ax.set_ylim(-10, 10)
leg = ax.legend(frameon=True)
leg.get_frame().set_edgecolor('black')
leg.get_frame().set_linewidth(2.0)
```

```python

```

notes/Test code HAL potential and effective masses.ipynb