# Assignment 4: Chatbot

Welcome to the last assignment of Course 4. Before you get started, we want to congratulate you on getting here. It is your 16th programming assignment in this Specialization and we are very proud of you! In this assignment, you are going to use the [Reformer](https://arxiv.org/abs/2001.04451), also known as the efficient Transformer, to generate a dialogue between two bots. You will feed conversations to your model and it will learn how to understand the context of each one. Not only will it learn how to answer questions but it will also know how to ask questions if it needs more info. For example, after a customer asks for a train ticket, the chatbot can ask what time the said customer wants to leave. You can use this concept to automate call centers, hotel receptions, personal trainers, or any type of customer service. By completing this assignment, you will:

* Understand how the Reformer works
* Explore the [MultiWoz](https://arxiv.org/abs/1810.00278) dataset
* Process the data to feed it into the model
* Train your model
* Generate a dialogue by feeding a question to the model

## Outline
- [Part 1:   Exploring the MultiWoz dataset](#1)
	- [Exercise 01](#ex01)
- [Part 2:   Processing the data for Reformer inputs](#2)
    - [2.1   Tokenizing, batching with bucketing](#2.1)
- [Part 3:   Reversible layers](#3)
	- [Exercise 02](#ex02)
	- [Exercise 03](#ex03)
    - [3.1   Reversible layers and randomness](#3.1)
- [Part 4:   ReformerLM Training](#4)
	- [Exercise 04](#ex04)
	- [Exercise 05](#ex05)
- [Part 5:   Decode from a pretrained model](#5)
	- [Exercise 06](#ex06)

# Part 1:   Exploring the MultiWoz dataset

You will start by exploring the MultiWoz dataset. The dataset you are about to use has more than 10,000 human annotated dialogues and spans multiple domains and topics. Some dialogues include multiple domains and others include single domains. In this section, you will load and explore this dataset, as well as develop a function to extract the dialogues.

Let's first import the modules we will be using:

```python
import json
import random
import numpy as np
from termcolor import colored

import trax
from trax import layers as tl
from trax.supervised import training
!pip list | grep trax
```

    INFO:tensorflow:tokens_length=568 inputs_length=512 targets_length=114 noise_density=0.15 mean_noise_span_length=3.0
    trax                     1.3.4
    WARNING: You are using pip version 20.1.1; however, version 20.2.3 is available.
    You should consider upgrading via the '/opt/conda/bin/python3 -m pip install --upgrade pip' command.

Let's also declare some constants we will be using in the exercises.

```python
# filename of the MultiWOZ dialogue dataset
DATA_FILE = 'data.json'

# data directory
DATA_DIR = './data'

# dictionary where we will load the dialogue dataset
DIALOGUE_DB = {}

# vocabulary filename
VOCAB_FILE = 'en_32k.subword'

# vocabulary file directory
VOCAB_DIR = 'data/vocabs'
```

Let's now load the MultiWOZ 2.1 dataset. We have already provided it for you in your workspace. It is in JSON format so we should load it as such:

```python
# help function to load a JSON file
def load_json(directory, file):
    with open(f'{directory}/{file}') as file:
        db = json.load(file)
    return db

# load the dialogue data set into our dictionary
DIALOGUE_DB = load_json(DATA_DIR, DATA_FILE)
```

Let's see how many dialogues we have in the dictionary. 1 key-value pair is one dialogue so we can just get the dictionary's length.

```python
print(f'The number of dialogues is: {len(DIALOGUE_DB)}')
```

    The number of dialogues is: 10438

The dialogues are composed of multiple files and the filenames are used as keys in our dictionary. Those with multi-domain dialogues have "MUL" in their filenames while single domain dialogues have either "SNG" or "WOZ".

```python
# print 7 keys from the dataset to see the filenames
print(list(DIALOGUE_DB.keys())[0:7])
```

    ['SNG01856.json', 'SNG0129.json', 'PMUL1635.json', 'MUL2168.json', 'SNG0073.json', 'SNG01445.json', 'MUL2105.json']

As you can see from the cells above, there are 10,438 conversations, each in its own file.  You will train your model on all those conversations. Each file is also loaded into a dictionary and each has two keys which are the following:

```python
# get keys of the fifth file in the list above
print(DIALOGUE_DB['SNG0073.json'].keys())
```

    dict_keys(['goal', 'log'])

The `goal` also points to a dictionary and it contains several keys pertaining to the objectives of the conversation. For example below, we can see that the conversation will be about booking a taxi.

```python
DIALOGUE_DB['SNG0073.json']['goal']
```

    {'taxi': {'info': {'leaveAt': '17:15',
       'destination': 'pizza hut fen ditton',
       'departure': "saint john's college"},
      'reqt': ['car type', 'phone'],
      'fail_info': {}},
     'police': {},
     'hospital': {},
     'hotel': {},
     'attraction': {},
     'train': {},
     'message': ["You want to book a taxi. The taxi should go to pizza hut fen ditton and should depart from saint john's college",
      "The taxi should leave after 17:15",
      "Make sure you get car type and contact number"],
     'restaurant': {}}

The `log` on the other hand contains the dialog. It is a list of dictionaries and each element of this list contains several descriptions as well. Let's look at an example:

```python
# get first element of the log list
DIALOGUE_DB['SNG0073.json']['log'][0]
```

    {'text': "I would like a taxi from Saint John's college to Pizza Hut Fen Ditton.",
     'metadata': {},
     'dialog_act': {'Taxi-Inform': [['Dest', 'pizza hut fen ditton'],
       ['Depart', "saint john 's college"]]},
     'span_info': [['Taxi-Inform', 'Dest', 'pizza hut fen ditton', 11, 14],
      ['Taxi-Inform', 'Depart', "saint john 's college", 6, 9]]}

For this assignment, we are only interested in the conversation which is in the `text` field.
The conversation goes back and forth between two persons. Let's call them 'Person 1' and 'Person 2'. This implies that
data['SNG0073.json']['log'][0]['text'] is 'Person 1' and
data['SNG0073.json']['log'][1]['text'] is 'Person 2' and so on. The even offsets are 'Person 1' and the odd offsets are 'Person 2'.

```python
print(' Person 1: ', DIALOGUE_DB['SNG0073.json']['log'][0]['text'])
print(' Person 2: ',DIALOGUE_DB['SNG0073.json']['log'][1]['text'])
```

     Person 1:  I would like a taxi from Saint John's college to Pizza Hut Fen Ditton.
     Person 2:  What time do you want to leave and what time do you want to arrive by?

### Exercise 01

You will now implement the `get_conversation()` function that will extract the conversations from the dataset's file.

**Instructions:** Implement a function to extract conversations from the input file.
As described above, the conversation is in the `text` field in each of the elements in the `log` list of the file. If the log list has `x` number of elements, then the function will get the `text` entries of each of those elements. Your function should return the conversation, prepending each field with either ' Person 1: ' if 'x' is even or ' Person 2: ' if 'x' is odd. You can use the Python modulus operator '%' to help select the even/odd entries. Important note: Do not print a newline character (i.e. `\n`) when generating the string. For example, in the code cell above, your function should output something like:

```
 Person 1: I would like a taxi from Saint John's college to Pizza Hut Fen Ditton. Person 2: What time do you want to leave and what time do you want to arrive by?
```

and **not**:

```
 Person 1:  I would like a taxi from Saint John's college to Pizza Hut Fen Ditton.
 Person 2:  What time do you want to leave and what time do you want to arrive by?
```

```python
# UNQ_C1
# GRADED FUNCTION: get_conversation
def get_conversation(file, data_db):
    '''
    Args:
        file (string): filename of the dialogue file saved as json
        data_db (dict): dialogue database

    Returns:
        string: A string containing the 'text' fields of  data[file]['log'][x]
    '''

    # initialize empty string
    result = ''

    # get length of file's log list
    len_msg_log = len(data_db[file]['log'])

    # set the delimiter strings
    delimiter_1 = ' Person 1: '
    delimiter_2 = ' Person 2: '

    # loop over the file's log list
    for i in range(len_msg_log):

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

        # get i'th element of file log list
        cur_log = data_db[file]['log'][i]

        # check if i is even
        if i%2 == 0:
            # append the 1st delimiter string
            result += delimiter_1
        else:
            # append the 2nd delimiter string
            result += delimiter_2

        # append the message text from the log
        result += cur_log['text']

    ### END CODE HERE ###

    return result
```

```python
# BEGIN UNIT TEST
import w4_unittest
w4_unittest.test_get_conversation(get_conversation)
# END UNIT TEST
```

     All tests passed

```python
file = 'SNG01856.json'
conversation = get_conversation(file, DIALOGUE_DB)

# print raw output
print(conversation)
```

     Person 1: am looking for a place to to stay that has cheap price range it should be in a type of hotel Person 2: Okay, do you have a specific area you want to stay in? Person 1: no, i just need to make sure it's cheap. oh, and i need parking Person 2: I found 1 cheap hotel for you that includes parking. Do you like me to book it? Person 1: Yes, please. 6 people 3 nights starting on tuesday. Person 2: I am sorry but I wasn't able to book that for you for Tuesday. Is there another day you would like to stay or perhaps a shorter stay? Person 1: how about only 2 nights. Person 2: Booking was successful.
    Reference number is : 7GAWK763. Anything else I can do for you? Person 1: No, that will be all. Good bye. Person 2: Thank you for using our services.

**Expected Result:**
```
Person 1: am looking for a place to to stay that has cheap price range it should be in a type of hotel Person 2: Okay, do you have a specific area you want to stay in? Person 1: no, i just need to make sure it's cheap. oh, and i need parking Person 2: I found 1 cheap hotel for you that includes parking. Do you like me to book it? Person 1: Yes, please. 6 people 3 nights starting on tuesday. Person 2: I am sorry but I wasn't able to book that for you for Tuesday. Is there another day you would like to stay or perhaps a shorter stay? Person 1: how about only 2 nights. Person 2: Booking was successful.
Reference number is : 7GAWK763. Anything else I can do for you? Person 1: No, that will be all. Good bye. Person 2: Thank you for using our services.
```

We can have a utility pretty print function just so we can visually follow the conversation more easily.

```python
def print_conversation(conversation):

    delimiter_1 = 'Person 1: '
    delimiter_2 = 'Person 2: '

    split_list_d1 = conversation.split(delimiter_1)

    for sublist in split_list_d1[1:]:
        split_list_d2 = sublist.split(delimiter_2)
        print(colored(f'Person 1: {split_list_d2[0]}', 'red'))

        if len(split_list_d2) > 1:
            print(colored(f'Person 2: {split_list_d2[1]}', 'green'))

print_conversation(conversation)
```

    Person 1: am looking for a place to to stay that has cheap price range it should be in a type of hotel
    Person 2: Okay, do you have a specific area you want to stay in?
    Person 1: no, i just need to make sure it's cheap. oh, and i need parking
    Person 2: I found 1 cheap hotel for you that includes parking. Do you like me to book it?
    Person 1: Yes, please. 6 people 3 nights starting on tuesday.
    Person 2: I am sorry but I wasn't able to book that for you for Tuesday. Is there another day you would like to stay or perhaps a shorter stay?
    Person 1: how about only 2 nights.
    Person 2: Booking was successful.
    Reference number is : 7GAWK763. Anything else I can do for you?
    Person 1: No, that will be all. Good bye.
    Person 2: Thank you for using our services.

For this assignment, we will just use the outputs of the calls to `get_conversation` to train the model. But just to expound, there are also other information in the MultiWoz dataset that can be useful in other contexts. Each element of the log list has more information about it. For example, above, if you were to look at the other fields for the following, "am looking for a place to stay that has cheap price range it should be in a type of hotel", you will get the following.

```python
DIALOGUE_DB['SNG01856.json']['log'][0]
```

    {'text': 'am looking for a place to to stay that has cheap price range it should be in a type of hotel',
     'metadata': {},
     'dialog_act': {'Hotel-Inform': [['Type', 'hotel'], ['Price', 'cheap']]},
     'span_info': [['Hotel-Inform', 'Type', 'hotel', 20, 20],
      ['Hotel-Inform', 'Price', 'cheap', 10, 10]]}

The dataset also comes with hotel, hospital, taxi, train, police, and restaurant databases. For example, in case you need to call a doctor, or a hotel, or a taxi, this will allow you to automate the entire conversation. Take a look at the files accompanying the data set.

```python
# this is an example of the attractions file
attraction_file = open('data/attraction_db.json')
attractions = json.load(attraction_file)
print(attractions[0])
```

    {'address': 'pool way, whitehill road, off newmarket road', 'area': 'east', 'entrance fee': '?', 'id': '1', 'location': [52.208789, 0.154883], 'name': 'abbey pool and astroturf pitch', 'openhours': '?', 'phone': '01223902088', 'postcode': 'cb58nt', 'pricerange': '?', 'type': 'swimmingpool'}

```python
# this is an example of the hospital file
hospital_file = open('data/hospital_db.json')
hospitals = json.load(hospital_file)
print(hospitals[0]) # feel free to index into other indices
```

    {'department': 'neurosciences critical care unit', 'id': 0, 'phone': '01223216297'}

```python
# this is an example of the hotel file
hotel_file = open('data/hotel_db.json')
hotels = json.load(hotel_file)
print(hotels[0]) # feel free to index into other indices
```

    {'address': '124 tenison road', 'area': 'east', 'internet': 'yes', 'parking': 'no', 'id': '0', 'location': [52.1963733, 0.1987426], 'name': 'a and b guest house', 'phone': '01223315702', 'postcode': 'cb12dp', 'price': {'double': '70', 'family': '90', 'single': '50'}, 'pricerange': 'moderate', 'stars': '4', 'takesbookings': 'yes', 'type': 'guesthouse'}

```python
# this is an example of the police file
police_file = open('data/police_db.json')
police = json.load(police_file)
print(police[0]) # feel free to index into other indices
```

    {'name': 'Parkside Police Station', 'address': 'Parkside, Cambridge', 'id': 0, 'phone': '01223358966'}

```python
# this is an example of a restuarant file
restaurant_file = open('data/restaurant_db.json')
restaurants = json.load(restaurant_file)
print(restaurants[0]) # feel free to index into other indices
```

    {'address': 'Regent Street City Centre', 'area': 'centre', 'food': 'italian', 'id': '19210', 'introduction': 'Pizza hut is a large chain with restaurants nationwide offering convenience pizzas pasta and salads to eat in or take away', 'location': [52.20103, 0.126023], 'name': 'pizza hut city centre', 'phone': '01223323737', 'postcode': 'cb21ab', 'pricerange': 'cheap', 'type': 'restaurant'}

For more information about the multiwoz 2.1 data set, please run the cell below to read the `ReadMe.txt` file. Feel free to open any other file to explore it.

```python
with open('data/README') as file:
    print(file.read())
```

    #####################################################
    #####################################################
    #  Copyright Cambridge Dialogue Systems Group, 2018 #
    #####################################################
    #####################################################

    Dataset contains the following files:
    1. data.json: the woz dialogue dataset, which contains the conversation  users and wizards, as well as a set of coarse labels for each user turn. This file contains both system and user dialogue acts annotated at the turn level. Files with multi-domain dialogues have "MUL" in their names. Single domain dialogues have either "SNG" or "WOZ" in their names.
    2. restaurant_db.json: the Cambridge restaurant database file, containing restaurants in the Cambridge UK area and a set of attributes.
    3. attraction_db.json: the Cambridge attraction database file, contining attractions in the Cambridge UK area and a set of attributes.
    4. hotel_db.json: the Cambridge hotel database file, containing hotels in the Cambridge UK area and a set of attributes.
    5. train_db.json: the Cambridge train (with artificial connections) database file, containing trains in the Cambridge UK area and a set of attributes.
    6. hospital_db.json: the Cambridge hospital database file, contatining information about departments.
    7. police_db.json: the Cambridge police station information.
    8. taxi_db.json: slot-value list for taxi domain.
    9. valListFile.txt: list of dialogues for validation.
    10. testListFile.txt: list of dialogues for testing.
    11. system_acts.json:
      There are 6 domains ('Booking', 'Restaurant', 'Hotel', 'Attraction', 'Taxi', 'Train') and 1 dummy domain ('general').
      A domain-dependent dialogue act is defined as a domain token followed by a domain-independent dialogue act, e.g. 'Hotel-inform' means it is an 'inform' act in the Hotel domain.
      Dialogue acts which cannot take slots, e.g., 'good bye', are defined under the 'general' domain.
      A slot-value pair defined as a list with two elements. The first element is slot token and the second one is its value.
      If a dialogue act takes no slots, e.g., dialogue act 'offer booking' for an utterance 'would you like to take a reservation?', its slot-value pair is ['none', 'none']
      There are four types of values:
      1) If a slot takes a binary value, e.g., 'has Internet' or 'has park', the value is either 'yes' or 'no'.
      2) If a slot is under the act 'request', e.g., 'request' about 'area', the value is expressed as '?'.
      3) The value that appears in the utterance e.g., the name of a restaurant.
      4) If for some reason the turn does not have an annotation then it is labeled as "No Annotation."
    12. ontology.json: Data-based ontology containing all the values for the different slots in the domains.
    13. slot_descriptions.json: A collection of human-written slot descriptions for each slot in the dataset. Each slot has at least two descriptions.
    14. tokenization.md: A description of the tokenization preprocessing we had to perform to maintain consistency between the dialogue act annotations of DSTC 8 Track 1 and the existing MultiWOZ 2.0 data.

As you can see, there are many other aspects of the MultiWoz dataset. Nonetheless, you'll see that even with just the conversations, your model will still be able to generate useful responses. This concludes our exploration of the dataset. In the next section, we will do some preprocessing before we feed it into our model for training.

# Part 2:   Processing the data for Reformer inputs

You will now use the `get_conversation()` function to process the data. The Reformer expects inputs of this form:

**Person 1: Why am I so happy? Person 2: Because you are learning NLP Person 1: ... Person 2: ...***

And the conversation keeps going with some text. As you can see 'Person 1' and 'Person 2' act as delimiters so the model automatically recognizes the person and who is talking. It can then come up with the corresponding text responses for each person. Let's proceed to process the text in this fashion for the Reformer. First, let's grab all the conversation strings from all dialogue files and put them in a list.

```python
# the keys are the file names
all_files = DIALOGUE_DB.keys()

# initialize empty list
untokenized_data = []

# loop over all files
for file in all_files:
    # this is the graded function you coded
    # returns a string delimited by Person 1 and Person 2
    result = get_conversation(file, DIALOGUE_DB)

    # append to the list
    untokenized_data.append(result)

# print the first element to check if it's the same as the one we got before
print(untokenized_data[0])
```

     Person 1: am looking for a place to to stay that has cheap price range it should be in a type of hotel Person 2: Okay, do you have a specific area you want to stay in? Person 1: no, i just need to make sure it's cheap. oh, and i need parking Person 2: I found 1 cheap hotel for you that includes parking. Do you like me to book it? Person 1: Yes, please. 6 people 3 nights starting on tuesday. Person 2: I am sorry but I wasn't able to book that for you for Tuesday. Is there another day you would like to stay or perhaps a shorter stay? Person 1: how about only 2 nights. Person 2: Booking was successful.
    Reference number is : 7GAWK763. Anything else I can do for you? Person 1: No, that will be all. Good bye. Person 2: Thank you for using our services.

Now let us split the list to a train and eval dataset.

```python
# shuffle the list we generated above
random.shuffle(untokenized_data)

# define a cutoff (5% of the total length for this assignment)
# convert to int because we will use it as a list index
cut_off = int(len(untokenized_data) * .05)

# slice the list. the last elements after the cut_off value will be the eval set. the rest is for training.
train_data, eval_data = untokenized_data[:-cut_off], untokenized_data[-cut_off:]

print(f'number of conversations in the data set: {len(untokenized_data)}')
print(f'number of conversations in train set: {len(train_data)}')
print(f'number of conversations in eval set: {len(eval_data)}')
```

    number of conversations in the data set: 10438
    number of conversations in train set: 9917
    number of conversations in eval set: 521

## 2.1   Tokenizing, batching with bucketing
We can now proceed in generating tokenized batches of our data. Let's first define a utility generator function to yield elements from our data sets:

```python
def stream(data):
    # loop over the entire data
    while True:
        # get a random element
        d = random.choice(data)

        # yield a tuple pair of identical values
        # (i.e. our inputs to the model will also be our targets during training)
        yield (d, d)
```

Now let's define our data pipeline for tokenizing and batching our data. As in the previous assignments, we will bucket by length and also have an upper bound on the token length.

```python
# trax allows us to use combinators to generate our data pipeline
data_pipeline = trax.data.Serial(
    # randomize the stream
    trax.data.Shuffle(),

    # tokenize the data
    trax.data.Tokenize(vocab_dir=VOCAB_DIR,
                       vocab_file=VOCAB_FILE),

    # filter too long sequences
    trax.data.FilterByLength(2048),

    # bucket by length
    trax.data.BucketByLength(boundaries=[128, 256,  512, 1024],
                             batch_sizes=[16,    8,    4,   2, 1]),

    # add loss weights but do not add it to the padding tokens (i.e. 0)
    trax.data.AddLossWeights(id_to_mask=0)
)

# apply the data pipeline to our train and eval sets
train_stream = data_pipeline(stream(train_data))
eval_stream = data_pipeline(stream(eval_data))
```

Peek into the train stream.

```python
# the stream generators will yield (input, target, weights). let's just grab the input for inspection
inp, _, _ = next(train_stream)

# print the shape. format is (batch size, token length)
print("input shape: ", inp.shape)

# detokenize the first element
print(trax.data.detokenize(inp[0], vocab_dir=VOCAB_DIR, vocab_file=VOCAB_FILE))
```

    input shape:  (4, 512)
     Person 1: I love the centre of Cambridge but I always spend too much, are there cheap places to eat there so I can stay within my budget? Person 2: Yes, there are several. Charlie Chan serves Chinese food, there's Mahal of Cambridge for Indian, La Raza has Spanish cuisine, and there are a few others. Would you like more? Person 1: restaurant should serve indian food is what I need  Person 2: There are three cheap Indian restaurants in the centre: Kohinoor, The Gandhi, and Mahal of Cambridge. Would you like to make a reservation at one? Person 1: HOw about mahal and can you give their phone number and postcode? Person 2: Their phone number is 01223360409 and their postal code is cb21rq. Person 1: Great.  I'm also looking for a train.  I need to leave after 16:15 from London Liverpool Street. Person 2: Which day would you like to travel by train? Person 1: I'll need it to be on Tuesday. I'm going to Cambridge.  Person 2: TR2855 will leave at 17:39, would that work for you? Person 1: make a booking for 1 people.
    Make sure you get the reference number Person 2: I have booked the ticket for you, reference 4O9ZT5ZD. Person 1: thank you very much. Person 2: Is there anything else I can assist you with today?  Person 1: That is all, make it a great day! Person 2: Thanks! You have a great day as well!

# Part 3:   Reversible layers

When running large deep models, you will often run out of memory as each layer allocates memory to store activations for use in backpropagation. To save this resource, you need to be able to recompute these activations during the backward pass without storing them during the forward pass. Take a look first at the leftmost diagram below.

This is how the residual networks are implemented in the standard Transformer. It follows that, given `F()` is Attention and `G()` is Feed-forward(FF).
:

\begin{align}
\mathrm{y}_\mathrm{a} &= \mathrm{x} + \mathrm{F}\left(\mathrm{x}\right)\tag{1} \\
\mathrm{y}_{b}&=\mathrm{y}_{a}+\mathrm{G}\left(\mathrm{y}_{a}\right)\tag{2}\\
\end{align}

As you can see, it requires that $\mathrm{x}$ and $\mathrm{y}_{a}$ be saved so it can be used during backpropagation. We want to avoid this to conserve memory and this is where reversible residual connections come in. They are shown in the middle and rightmost diagrams above. The key idea is that we will start with two copies of the input to the model and at each layer we will only update one of them. The activations that we *don’t* update are the ones that will be used to compute the residuals.

Now in this reversible set up you get the following instead:

\begin{align}
\mathrm{y}_{1}&=\mathrm{x}_{1}+\mathrm{F}\left(\mathrm{x}_{2}\right)\tag{3}\\
\mathrm{y}_{2}&=\mathrm{x}_{2}+\mathrm{G}\left(\mathrm{y}_{1}\right)\tag{4}\\
\end{align}
To recover $\mathrm{(x_1,x_2)}$ from $\mathrm{(y_1, y_2)}$

\begin{align}
\mathrm{x}_{2}&=\mathrm{y}_{2}-\mathrm{G}\left(\mathrm{y}_{1}\right)\tag{5}\\
\mathrm{x}_{1}&=\mathrm{y}_{1}-\mathrm{F}\left(\mathrm{x}_{2}\right)\tag{6}\\
\end{align}

With this configuration, we’re now able to run the network fully in reverse. You'll notice that during the backward pass, $\mathrm{x2}$ and $\mathrm{x1}$ can be recomputed based solely on the values of $\mathrm{y2}$ and $\mathrm{y1}$. No need to save it during the forward pass.

### Exercise 02
**Instructions:** You will implement the `reversible_layer_forward` function using equations 3 and 4 above. This function takes in the input vector `x` and the functions `f` and `g` and returns the concatenation of $y_1 and y_2$. For this exercise, we will be splitting `x` before going through the reversible residual steps$\mathrm{^1}$.  We can then use those two vectors for the `reversible_layer_reverse` function. Utilize `np.concatenate()` to form the output being careful to match the axis of the `np.split()`.

$\mathrm{^1}$*Take note that this is just for demonstrating the concept in this exercise and there are other ways of processing the input. As you'll see in the Reformer architecture later, the initial input (i.e. `x`) can instead be duplicated instead of split.*

```python
# UNQ_C2
# GRADED FUNCTION: reversible_layer_forward
def reversible_layer_forward(x, f, g):
    """
    Args:
        x (np.array): an input vector or matrix
        f (function): a function which operates on a vector/matrix
        g (function): a function which operates on a vector/matrix
    Returns:
        y (np.array): an output vector or matrix whose form is determined by 'x', f and g
    """
    # split the input vector into two (* along the last axis because it is the depth dimension)
    x1, x2 = np.split(x, 2, axis=-1)

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

    # get y1 using equation 3
    y1 = x1 + f(x2)

    # get y2 using equation 4
    y2 = x2 + g(y1)

    # concatenate y1 and y2 along the depth dimension. be sure output is of type np.ndarray
    y = np.concatenate([y1, y2], axis=-1)

    ### END CODE HERE ###
    return y
```

```python
# BEGIN UNIT TEST
w4_unittest.test_reversible_layer_forward(reversible_layer_forward)
# END UNIT TEST
```

     All tests passed

### Exercise 03

You will now implement the `reversible_layer_reverse` function  which is possible because at every time step you have $x_1$ and $x_2$ and $y_2$ and $y_1$, along with the function `f`, and `g`. Where `f` is the attention and `g` is the feedforward. This allows you to compute equations 5 and 6.

**Instructions:** Implement the `reversible_layer_reverse`. Your function takes in the output vector from  `reversible_layer_forward` and functions f and g. Using equations 5 and 6 above, it computes the inputs to the layer,  $x_1$ and $x_2$.  The output, x, is the concatenation of  $x_1, x_2$. Utilize `np.concatenate()`  to form the output being careful to match the axis of the `np.split()`.

```python
# UNQ_C3
# GRADED FUNCTION: reversible_layer_reverse
def reversible_layer_reverse(y, f, g):
    """
    Args:
        y (np.array): an input vector or matrix
        f (function): a function which operates on a vector/matrix of the form of 'y'
        g (function): a function which operates on a vector/matrix of the form of 'y'
    Returns:
        y (np.array): an output vector or matrix whose form is determined by 'y', f and g
    """

    # split the input vector into two (* along the last axis because it is the depth dimension)
    y1, y2 = np.split(y, 2, axis=-1)

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

    # compute x2 using equation 5
    x2 = y2 - g(y1)

    # compute x1 using equation 6
    x1 = y1 - f(x2)

    # concatenate x1 and x2 along the depth dimension
    x = np.concatenate([x1, x2], axis=-1)

    ### END CODE HERE ###
    return x

```

```python
# BEGIN UNIT TEST
w4_unittest.test_reversible_layer_reverse(reversible_layer_reverse)
# END UNIT TEST
```

     All tests passed

```python
# UNIT TEST COMMENT: assert at the end can be used in grading as well
f = lambda x: x + 2
g = lambda x: x * 3
input_vector = np.random.uniform(size=(32,))

output_vector = reversible_layer_forward(input_vector, f, g)
reversed_vector = reversible_layer_reverse(output_vector, f, g)

assert np.allclose(reversed_vector, input_vector)
```

## 3.1   Reversible layers and randomness

This is why we were learning about fastmath's random functions and keys in Course 3 Week 1. Utilizing the same key, `trax.fastmath.random.uniform()` will return the same values. This is required for the backward pass to return the correct layer inputs when random noise is introduced in the layer.

```python
# Layers like dropout have noise, so let's simulate it here:
f = lambda x: x + np.random.uniform(size=x.shape)

# See that the above doesn't work any more:
output_vector = reversible_layer_forward(input_vector, f, g)
reversed_vector = reversible_layer_reverse(output_vector, f, g)

assert not np.allclose(reversed_vector, input_vector)  # Fails!!

# It failed because the noise when reversing used a different random seed.

random_seed = 27686
rng = trax.fastmath.random.get_prng(random_seed)
f = lambda x: x + trax.fastmath.random.uniform(key=rng, shape=x.shape)

# See that it works now as the same rng is used on forward and reverse.
output_vector = reversible_layer_forward(input_vector, f, g)
reversed_vector = reversible_layer_reverse(output_vector, f, g)

assert np.allclose(reversed_vector, input_vector,  atol=1e-07)
```

# Part 4:   ReformerLM Training

You will now proceed to training your model. Since you have already know the two main components that differentiates it from the standard Transformer, LSH in Course 1 and reversible layers above, you can just use the pre-built model already implemented in Trax. It will have this architecture:

Similar to the Transformer you learned earlier, you want to apply an attention and feed forward layer to your inputs. For the Reformer, we improve the memory efficiency by using **reversible decoder blocks** and you can picture its implementation in Trax like below:

You can see that it takes the initial inputs `x1` and `x2` and does the first equation of the reversible networks you learned in Part 3. As you've also learned, the reversible residual has two equations for the forward-pass so doing just one of them will just constitute half of the reversible decoder block. Before doing the second equation (i.e. second half of the reversible residual), it first needs to swap the elements to take into account the stack semantics in Trax. It simply puts `x2` on top of the stack so it can be fed to the add block of the half-residual layer. It then swaps the two outputs again so it can be fed to the next layer of the network. All of these arrives at the two equations in Part 3 and it can be used to recompute the activations during the backward pass.

These are already implemented for you in Trax and in the following exercise, you'll get to practice how to call them to build your network.

### Exercise 04
**Instructions:** Implement a wrapper function that returns a Reformer Language Model. You can use Trax's [ReformerLM](https://trax-ml.readthedocs.io/en/latest/trax.models.html#trax.models.reformer.reformer.ReformerLM) to do this quickly. It will have the same architecture as shown above.

```python
# UNQ_C4
# GRADED FUNCTION
def ReformerLM(vocab_size=33000, n_layers=2, mode='train', attention_type=tl.SelfAttention):
    """
    Args:
        vocab_size (int): size of the vocabulary
        n_layers (int): number of decoder layers
        mode (string): setting of the model which can be 'train', 'eval', or 'predict'
        attention_type(class): attention class to use
    Returns:
        model (ReformerLM): a reformer language model implemented in Trax
    """

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###
    # initialize an instance of Trax's ReformerLM class
    model = trax.models.reformer.ReformerLM(
        # set vocab size
        vocab_size=vocab_size,
        # set number of layers
        n_layers=n_layers,
        # set mode
        mode=mode,
        # set attention type
        attention_type=attention_type
    )

    ### END CODE HERE ###
    return model
```

```python
# display the model
temp_model = ReformerLM('train')
print(str(temp_model))

# free memory
del temp_model
```

    Serial[
      ShiftRight(1)
      Embedding_train_512
      Dropout
      PositionalEncoding
      Dup_out2
      ReversibleSerial_in2_out2[
        ReversibleHalfResidualV2_in2_out2[
          Serial[
            LayerNorm
          ]
          SelfAttention
        ]
        ReversibleSwap_in2_out2
        ReversibleHalfResidualV2_in2_out2[
          Serial[
            LayerNorm
            Dense_2048
            Dropout
            FastGelu
            Dense_512
            Dropout
          ]
        ]
        ReversibleSwap_in2_out2
        ReversibleHalfResidualV2_in2_out2[
          Serial[
            LayerNorm
          ]
          SelfAttention
        ]
        ReversibleSwap_in2_out2
        ReversibleHalfResidualV2_in2_out2[
          Serial[
            LayerNorm
            Dense_2048
            Dropout
            FastGelu
            Dense_512
            Dropout
          ]
        ]
        ReversibleSwap_in2_out2
      ]
      Concatenate_in2
      LayerNorm
      Dropout
      Dense_train
      LogSoftmax
    ]

```python
# BEGIN UNIT TEST
w4_unittest.test_ReformerLM(ReformerLM)
# END UNIT TEST
```

     All tests passed

### Exercise 05
You will now write a function that takes in your model and trains it.

**Instructions:** Implement the `training_loop` below to train the neural network above. Here is a list of things you should do:

- Create `TrainTask` and `EvalTask`
- Create the training loop `trax.supervised.training.Loop`
- Pass in the following depending to train_task :
    - `labeled_data=train_gen`
    - `loss_layer=tl.CrossEntropyLoss()`
    - `optimizer=trax.optimizers.Adam(0.01)`
    - `lr_schedule=lr_schedule`
    - `n_steps_per_checkpoint=10`

You will be using your CrossEntropyLoss loss function with Adam optimizer. Please read the [trax](https://trax-ml.readthedocs.io/en/latest/trax.optimizers.html?highlight=adam#trax.optimizers.adam.Adam) documentation to get a full understanding.

- Pass in the following to eval_task:
     - `labeled_data=eval_gen`
     - `metrics=[tl.CrossEntropyLoss(), tl.Accuracy()]`

This function should return a `training.Loop` object. To read more about this check the [docs](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html?highlight=loop#trax.supervised.training.Loop).

```python
# UNQ_C5
# GRADED FUNCTION: train_model
def training_loop(ReformerLM, train_gen, eval_gen, output_dir = "./model/"):
    """
    Args:
        ReformerLM:  the Reformer language model you are building
        train_gen (generator): train data generator.
        eval_gen (generator): Validation generator.
        output_dir (string): Path to save the model output. Defaults to './model/'.

    Returns:
        trax.supervised.training.Loop: Training loop for the model.
    """

    # use the warmup_and_rsqrt_decay learning rate schedule
    lr_schedule = trax.lr.warmup_and_rsqrt_decay(
        n_warmup_steps=1000, max_value=0.01)

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

    # define the train task
    train_task = training.TrainTask(
        # labeled data
        labeled_data=train_gen,
        # loss layer
        loss_layer=tl.CrossEntropyLoss(),
        # optimizer
        optimizer=trax.optimizers.Adam(0.01),
        # lr_schedule
        lr_schedule=lr_schedule,
        # n_steps
        n_steps_per_checkpoint=10
    )

    # define the eval task
    eval_task = training.EvalTask(
        # labeled data
        labeled_data=eval_gen,
        # metrics
        metrics=[tl.CrossEntropyLoss(), tl.Accuracy()]
    )

    ### END CODE HERE ###
    loop = training.Loop(ReformerLM(mode='train'),
                         train_task,
                         eval_tasks=[eval_task],
                         output_dir=output_dir)
    return loop
```

```python
# UNIT TEST COMMENT: Use the train task and eval task for grading train_model
test_loop = training_loop(ReformerLM, train_stream, eval_stream)
train_task = test_loop._task
eval_task = test_loop._eval_task

print(train_task)
print(eval_task)
```

    <trax.supervised.training.TrainTask object at 0x7fa279760050>
    <trax.supervised.training.EvalTask object at 0x7fa279aba550>

```python
# BEGIN UNIT TEST
w4_unittest.test_tasks(train_task, eval_task)
# END UNIT TEST
```

     All tests passed

```python
# we will now test your function
!rm -f model/model.pkl.gz
loop = training_loop(ReformerLM, train_stream, eval_stream)
loop.run(10)
```

    Step      1: Ran 1 train steps in 58.85 secs
    Step      1: train CrossEntropyLoss |  10.44030952
    Step      1: eval  CrossEntropyLoss |  10.42661381
    Step      1: eval          Accuracy |  0.00000000

    Step     10: Ran 9 train steps in 133.20 secs
    Step     10: train CrossEntropyLoss |  10.27762985
    Step     10: eval  CrossEntropyLoss |  9.91913414
    Step     10: eval          Accuracy |  0.12019914

**Approximate Expected output:**

```

Step      1: Ran 1 train steps in 55.73 secs
Step      1: train CrossEntropyLoss |  10.41907787
Step      1: eval  CrossEntropyLoss |  10.41005802
Step      1: eval          Accuracy |  0.00000000

Step     10: Ran 9 train steps in 108.21 secs
Step     10: train CrossEntropyLoss |  10.15449715
Step     10: eval  CrossEntropyLoss |  9.63478279
Step     10: eval          Accuracy |  0.16350447
```

# Part 5:   Decode from a pretrained model

We will now proceed on decoding using the model architecture you just implemented. As in the previous weeks, we will be giving you a pretrained model so you can observe meaningful output during inference. You will be using the [autoregressive_sample_stream()](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.decoding.autoregressive_sample_stream) decoding method from Trax to do fast inference. Let's define a few parameters to initialize our model.

```python
# define the `predict_mem_len` and `predict_drop_len` of tl.SelfAttention
def attention(*args, **kwargs):
    # number of input positions to remember in a cache when doing fast inference.
    kwargs['predict_mem_len'] = 120
    # number of input elements to drop once the fast inference input cache fills up.
    kwargs['predict_drop_len'] = 120
    # return the attention layer with the parameters defined above
    return tl.SelfAttention(*args, **kwargs)

# define the model using the ReformerLM function you implemented earlier.
model = ReformerLM(
    vocab_size=33000,
    n_layers=6,
    mode='predict',
    attention_type=attention,
)

# define an input signature so we can initialize our model. shape will be (1, 1) and the data type is int32.
shape11 = trax.shapes.ShapeDtype((1, 1), dtype=np.int32)
```

We can now initialize our model from a file containing the pretrained weights. We will save this starting state so we can reset the model state when we generate a new conversation. This will become clearer in the `generate_dialogue()` function later.

```python
# initialize from file
model.init_from_file('chatbot_model1.pkl.gz',
                     weights_only=True, input_signature=shape11)

# save the starting state
STARTING_STATE = model.state
```

Let's define a few utility functions as well to help us tokenize and detokenize. We can use the [tokenize()](https://trax-ml.readthedocs.io/en/latest/trax.data.html#trax.data.tf_inputs.tokenize) and [detokenize()](https://trax-ml.readthedocs.io/en/latest/trax.data.html#trax.data.tf_inputs.detokenize) from `trax.data.tf_inputs` to do this.

```python
def tokenize(sentence, vocab_file, vocab_dir):
    return list(trax.data.tokenize(iter([sentence]), vocab_file=vocab_file, vocab_dir=vocab_dir))[0]

def detokenize(tokens, vocab_file, vocab_dir):
    return trax.data.detokenize(tokens, vocab_file=vocab_file, vocab_dir=vocab_dir)
```

We are now ready to define our decoding function. This will return a generator that yields that next symbol output by the model. It will be able to predict the next words by just feeding it a starting sentence.

### Exercise 06
**Instructions:** Implement the function below to return a generator that predicts the next word of the conversation.

```python
# UNQ_C6
# GRADED FUNCTION
def ReformerLM_output_gen(ReformerLM, start_sentence, vocab_file, vocab_dir, temperature):
    """
    Args:
        ReformerLM:  the Reformer language model you just trained
        start_sentence (string): starting sentence of the conversation
        vocab_file (string): vocabulary filename
        vocab_dir (string): directory of the vocabulary file
        temperature (float): parameter for sampling ranging from 0.0 to 1.0.
            0.0: same as argmax, always pick the most probable token
            1.0: sampling from the distribution (can sometimes say random things)

    Returns:
        generator: yields the next symbol generated by the model
    """

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

    # Create input tokens using the the tokenize function
    input_tokens = tokenize(start_sentence, vocab_file=vocab_file, vocab_dir=vocab_dir)

    # Add batch dimension to array. Convert from (n,) to (x, n) where
    # x is the batch size. Default is 1. (hint: you can use np.expand_dims() with axis=0)
    input_tokens_with_batch = np.array(input_tokens)[None, :]

    # call the autoregressive_sample_stream function from trax
    output_gen = trax.supervised.decoding.autoregressive_sample_stream(
        # model
        ReformerLM,
        # inputs will be the tokens with batch dimension
        inputs=input_tokens_with_batch,
        # temperature
        temperature=temperature
    )

    ### END CODE HERE ###

    return output_gen
```

```python
# BEGIN UNIT TEST
import pickle

WEIGHTS_FROM_FILE = ()

with open('weights', 'rb') as file:
    WEIGHTS_FROM_FILE = pickle.load(file)

shape11 = trax.shapes.ShapeDtype((1, 1), dtype=np.int32)

def attention(*args, **kwargs):
    kwargs['predict_mem_len'] = 120
    kwargs['predict_drop_len'] = 120
    return tl.SelfAttention(*args, **kwargs)

test_model = ReformerLM(vocab_size=5, n_layers=1, mode='predict', attention_type=attention)

test_output_gen = ReformerLM_output_gen(test_model, "test", vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, temperature=0)

test_model.init_weights_and_state(shape11)

test_model.weights = WEIGHTS_FROM_FILE

output = []

for i in range(6):
    output.append(next(test_output_gen)[0])

print(output)

# free memory
del test_model
del WEIGHTS_FROM_FILE
del test_output_gen
# END UNIT TEST
```

    [1, 0, 4, 3, 0, 4]

***Expected value:***

[1, 0, 4, 3, 0, 4]

Great! Now you will be able to see the model in action. The utility function below will call the generator you just implemented and will just format the output to be easier to read.

```python
shape11 = trax.shapes.ShapeDtype((1, 1), dtype=np.int32)

def attention(*args, **kwargs):
    kwargs['predict_mem_len'] = 120  # max length for predictions
    kwargs['predict_drop_len'] = 120  # never drop old stuff
    return tl.SelfAttention(*args, **kwargs)

model = ReformerLM(
    vocab_size=33000,
    n_layers=6,
    mode='predict',
    attention_type=attention,
)
```

```python
model.init_from_file('chatbot_model1.pkl.gz',
                     weights_only=True, input_signature=shape11)

STARTING_STATE = model.state
```

```python
def generate_dialogue(ReformerLM, model_state, start_sentence, vocab_file, vocab_dir, max_len, temperature):
    """
    Args:
        ReformerLM:  the Reformer language model you just trained
        model_state (np.array): initial state of the model before decoding
        start_sentence (string): starting sentence of the conversation
        vocab_file (string): vocabulary filename
        vocab_dir (string): directory of the vocabulary file
        max_len (int): maximum number of tokens to generate
        temperature (float): parameter for sampling ranging from 0.0 to 1.0.
            0.0: same as argmax, always pick the most probable token
            1.0: sampling from the distribution (can sometimes say random things)

    Returns:
        generator: yields the next symbol generated by the model
    """

    # define the delimiters we used during training
    delimiter_1 = 'Person 1: '
    delimiter_2 = 'Person 2: '

    # initialize detokenized output
    sentence = ''

    # token counter
    counter = 0

    # output tokens. we insert a ': ' for formatting
    result = [tokenize(': ', vocab_file=vocab_file, vocab_dir=vocab_dir)]

    # reset the model state when starting a new dialogue
    ReformerLM.state = model_state

    # calls the output generator implemented earlier
    output = ReformerLM_output_gen(ReformerLM, start_sentence, vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, temperature=temperature)

    # print the starting sentence
    print(start_sentence.split(delimiter_2)[0].strip())

    # loop below yields the next tokens until max_len is reached. the if-elif is just for prettifying the output.
    for o in output:

        result.append(o)

        sentence = detokenize(np.concatenate(result, axis=0), vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR)

        if sentence.endswith(delimiter_1):
            sentence = sentence.split(delimiter_1)[0]
            print(f'{delimiter_2}{sentence}')
            sentence = ''
            result.clear()

        elif sentence.endswith(delimiter_2):
            sentence = sentence.split(delimiter_2)[0]
            print(f'{delimiter_1}{sentence}')
            sentence = ''
            result.clear()

        counter += 1

        if counter > max_len:
            break

```

We can now feed in different starting sentences and see how the model generates the dialogue. You can even input your own starting sentence. Just remember to ask a question that covers the topics in the Multiwoz dataset so you can generate a meaningful conversation.

```python
sample_sentence = ' Person 1: Are there theatres in town? Person 2: '
generate_dialogue(ReformerLM=model, model_state=STARTING_STATE, start_sentence=sample_sentence, vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, max_len=120, temperature=0.2)
```

    Person 1: Are there theatres in town?
    Person 2: : There are 4 theatres in town. Do you have a preference on area?
    Person 1: No, I don't care. Which one do you recommend?
    Person 2: I recommend the Mumford Theatre. It's located in the east at Anglia Ruskin Enterprise, east road. Would you like more information?
    Person 1: Yes, could I get the postcode and phone number please?
    Person 2: The phone number is 08451962320 and expensive. Which one would you prefer?

```python
sample_sentence = ' Person 1: Is there a hospital nearby? Person 2: '
generate_dialogue(ReformerLM=model, model_state=STARTING_STATE, start_sentence=sample_sentence, vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, max_len=120, temperature=0.2)
```

    Person 1: Is there a hospital nearby?
    Person 2: : Addensbrookes Hospital is located at Hills Rd, Cambridge, postcode CB20QQ. Do you need anything else?
    Person 1: No, that's all. Thank you!
    Person 2: You're welcome. Have a good day.m.Thanks for contacting the Cambridge TownInfo centre. Goodbye.
    Person 1: Thanks again. Goodbye.

```python
sample_sentence = ' Person 1: Can you book a taxi? Person 2: '
generate_dialogue(ReformerLM=model, model_state=STARTING_STATE, start_sentence=sample_sentence, vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, max_len=120, temperature=0.2)
```

    Person 1: Can you book a taxi?
    Person 2: : I sure can. When would you like to leave?
    Person 1: I need to leave after 21:45.
    Person 2: I'm sorry I have no trains leaving from that day. Could you change your request?
    Person 1: How about a college?
    Person 2: I have 13 colleges in the center of town. Downing College is in the center of town.

```python
sample_sentence = ' Person 1: What is the most popular food in Canada? Person 2: '
generate_dialogue(ReformerLM=model, model_state=STARTING_STATE, start_sentence=sample_sentence, vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, max_len=120, temperature=0.2)
```

    Person 1: What is the most popular food in Canada?
    Person 2: : There are no restaurants that match your request. Would you like to try a different cuisine?
    Person 1: How about Chinese food?
    Person 2: There are several Chinese restaurants in town. Do you have a price range in mind?
    Person 1: Yes, I'd like to eat Chinese food.
    Person 2: There are several Chinese restaurants in town. Do you have a price range in mind?
    Person 1: I would like to eat in the centre of town. Person:  i i i want

**Congratulations! You just wrapped up the final assignment of this course and the entire specialization!**


##### Copyright 2021 The Cirq Developers

```
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
```

```
# Fourier Checking Problem
```

    ">View on QuantumAI

    Run in Google Colab

    View source on GitHub

    Download notebook

```
try:
    import cirq
except:
    print("installing cirq...")
    !pip install --quiet cirq
    print("installed cirq.")
    import cirq

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

np.random.seed(2021)
np.set_printoptions(precision=3, suppress=True, linewidth=200)
```

# Introduction

In past decades, there is ample evidence suggesting that quantum computers can be exponentially more powerful in solving certain computational tasks than their classical couterparts. The *black-box* or *query* model, such as Grover’s search, Deutsch-Jozsa’s algorithm, etc., offers a concrete setting to show such exponential speedups. Normally, one provides "black-box access" to a function $f$, meaning that the quantum algorithm can apply a unitary
transformation that maps basis states of the form $|x, y \rangle$ to to basis states of the form $|x, y \oplus f(x)\rangle $ or $|x\rangle$ to $(-1)^{f(x)} |x\rangle$ if $f$ is Boolean. Then, a natural question is asked:

> What is the maximal possible separation between quantum and classical query complexities?

For example, could there be a function of $N$ bits with a quantum query
complexity of 1, but a classical randomized query complexity of $\Omega(\sqrt{N})$ or $\Omega(N)$?
Specifically, [Buhrman et al. from 2002](https://www.sciencedirect.com/science/article/pii/S030439750100144X) asked whether there is any
property of $N-$bit strings that exhibits a **“maximal”** separation: that is, one that requires $\Omega(N)$
queries to test classically, but only O (1) quantumly.

**Fourier Checking** is a problem that provides a separation between quantum  and classical computers -- $O(1)$ VS $\tilde{\Omega}(\sqrt{N})$, which can be proved as optimal. Currently it only has theoretical importance - but, as it falls into the category of small quantum algorithms, it can be used to demonstrate query complexity and oracle synthesis in Cirq.

Goal of this notebook is to introduce

1. What is Forrelation and the Fourier Checking problem and why we are interested in it?
2. What is bounded-error quantum polynomial time (BQP) and why does the Fourier Checking problem belong to it?
3. How to implement the Fourier Checking algorithm and an oracle function in Cirq?

We won't include the formal proofs and argument. However, we do give sketches of the derivation for intuition and encourage the reader to check the corresponding lemmas and theorems in the original paper.

# Preliminary

Before we present the Fourier Checking problem, three preliminary concepts, 1) BPP and BQP 2) Fourier Transform over $Z_2^n$ and 3) Forrelation, are introduced first.

## Bounded-error Probabilistic Polynomial time (BPP)

In computational complexity theory, [bounded-error probabilistic polynomial time (BPP)](https://en.wikipedia.org/wiki/BPP_(complexity)) is the class of decision problems solvable by a [probabilistic Turing machine](https://en.wikipedia.org/wiki/Probabilistic_Turing_machine) in polynomial time with an error probability bounded away from 1/3 for all instances:

|        Anwser Provided

 Correct Anwser| Yes

   | No

  |
|:---------------|------|-------|
| Yes           | $\geq$ 2/3 | $\leq$ 1/3 |
| No            | $\leq$ 1/3 | $\geq$ 2/3 |

The choice of 1/3 in the definition is arbitrary. It can be any constant between 0 and 1/2 (exclusive) and the set BPP will be unchanged.

To understand this better, let's consider a simple example.

**Problem**: Supposing we have a boolean function $f: \{0, 1\}^n \rightarrow\{-1, 1\}$. The function is drawn either from distribution $\mathcal{B}$ or $\mathcal{C}$. Under the distribution $\mathcal{B}$, the function $f$ is balanced, i.e., $\sum_{x} f(x) = 0$. Under the distribution $\mathcal{C}$, the function $f$ is constant, i.e., function is always equal to $+1$ or $-1$. The problem is to accept the function if it is drawn from $\mathcal{B}$ or reject it otherwise.

For example, consider the case that $n=2$. There are only four possible functions listed in following table

|x   | f_0(x)| f_1(x)| f_2(x)| f_3(x)|
|----|----:|----:|----:|----:|
| 0  | +1 | -1 | +1 | -1 |
| 1  | -1 | +1 | +1 | -1 |

Clearly, $f_0$ and $f_1$ belongs to $\mathcal{B}$ and $f_2$ and $f_3$ belongs to $\mathcal{C}$.

**Deterministic Algorithm**:  We evaluate the outputs of function of $2^{n-1}+1$ different inputs. If the results contain both $+1$ and $-1$ value, the function must be drawn from distribution $\mathcal{B}$. Otherwise, it must come from $\mathcal{C}$.

Remembering that the function is guaranteed to be either balanced or constant, not somewhere in between. So above algorithm is guaranteed to be always correct. However, the query complexity is $O(N)$, denoting $N=2^n$. Here we are interested in the query complexity instead of computation complexity. Namely, how many times we have to evaluate $f$ -- which you can imagine being a very costly function. If the bounded-error probability is acceptable for solving the problem, we can achieve better algorithm in terms of query complexity.

**Randomized Algorithm**: Randomly select $K$ different inputs and evaluate the corresponding outputs. If both $+1$ and $-1$ are observed, we accept it. Otherwise, it reject it.

The complexity of algorithm depends on the choice of $K$. For this problem, selecting fixed number is sufficient, i.e. $O(1)$ query complexity. It is worth to remark that $O(1)$ means no matter the fixed number is independent of the input size of function $N$. Consider the $K=2$ case, it is not hard to establish the following confusion matrix regardless of $N$:

| Randomized Algorithm (K=2):  | Accept  | Reject |
|---------------|---------|-------|
| Drawn from $\mathcal{B}$ | 1/2 | 1/2  |
| Drawn from $\mathcal{C}$ | 0 | 1 |

This is not sufficient to solve it. However, if we select $K>2$ entries, the  probability of correctness will boost. Let's use the code to exam it.

```
def gen_balanced_function(N):
    half_size = N // 2
    f = np.ones(N)
    flip_loc = np.random.permutation(N)[:half_size]
    f[flip_loc] = -1
    return f

def gen_constant_function(N):
    flip = np.random.random() > 0.5
    f = np.ones(N) if flip else -1 * np.ones(N)
    return f
```

```
def randomized_alg(f, sample_size):
    N = len(f)
    sample_index = np.random.choice(N, size=sample_size)
    if len(set(f[sample_index])) == 2:
        return "accept"
    return "reject"
```

```
N = 128
K = 3
samples = 1000

res = pd.DataFrame()
for _ in range(samples):
    if np.random.rand() > 0.5:
        f = gen_balanced_function(N)
        dist = "B"
    else:
        f = gen_constant_function(N)
        dist = "C"
    decision = randomized_alg(f, K)
    res = res.append({
        "Distribution": dist,
        "Decision": decision,
        "Count": 1
    }, ignore_index=True)
confusion = res.pivot_table(index="Distribution",
                            columns="Decision",
                            values="Count",
                            aggfunc="sum")
# Translate the counts into percentage
confusion.div(confusion.sum(axis=1), axis=0).apply(lambda x: round(x, 4) * 100)
```

Recall that the B in BPP means *bounded-error*. Actually, you can achieve arbitary small error rate under the same query complexity order. Say, you set an acceptable error rate $\epsilon$. Then the key is that we can run the algorithm multiple times. Repeat it as many times as you want until the error rate is lower than $\epsilon$. It is crucial to note that the error rate for this particular problem does not depend on the size of the input but only on the size of the sample and the reptitation. For this reason, in order to get to a bounded probability error, it is sufficient to just adjust the sample size and/or repetitions to a given constant - which means that the "query complexity" of the algorithm will stay $O(1)$. For example, let's run previous algorithms 3 times and make the final decision based on the majority of the decision of each term. You should verify that the probability of error indeed decreased and independent of $N$.

```
N = 128
K = 3
repetitions = 3
samples = 1000

res = pd.DataFrame()
for _ in range(samples):
    if np.random.rand() > 0.5:
        f = gen_balanced_function(N)
        dist = "B"
    else:
        f = gen_constant_function(N)
        dist = "C"
    accept_minus_reject_count = 0
    for _ in range(repetitions):
        decision = randomized_alg(f, K)
        accept_minus_reject_count += 1 if decision == "accept" else -1
    final_decision = "accept" if accept_minus_reject_count > 0 else "reject"
    res = res.append(
        {
            "Distribution": dist,
            "Decision": final_decision,
            "Count": 1
        }, ignore_index=True)
confusion = res.pivot_table(index="Distribution",
                            columns="Decision",
                            values="Count",
                            aggfunc="sum")
# Translate the counts into percentage
confusion.div(confusion.sum(axis=1), axis=0).apply(lambda x: round(x, 4) * 100)
```

If you try the algorithm with more repetitions, you will observe the error rate decreases rapidly. You can also try different values of $K, N, $ and/or repetitions to see how the confusion matrix changes according.

After you understand the concept of BPP, it is easy to understand [bounded-error quantum polynomial time (BQP)](https://en.wikipedia.org/wiki/BQP) now. BQP is the class of decision problems solvable by a quantum computer in polynomial time, with an error probability of at most 1/3 for all instances. It is the quantum analogue to the complexity class BPP. Actually, under quantum computer, previous problem can be solved through [Deutsch–Jozsa algorithm](https://en.wikipedia.org/wiki/Deutsch%E2%80%93Jozsa_algorithm) with a single query while guaranteed to be correct always. The Fourier Checking problem that will be introduced later is belong to BQP as well.

## Fourier Transform over $\mathbb{Z}^n_2$
In this colab, we are interested in the boolean function of the form $f : \{0, 1\}^n \rightarrow\{-1, 1\}$. In this case, the Fourier transform of $f$ over $\mathbb{Z}^n_2$ is defined as
$$
  \hat{f}(y) := \frac{1}{\sqrt{N}} \sum_{x\in\{0,1\}^n} (-1)^{x \cdot y} f
  (x).
$$
where $x \cdot y$ means the bit-wise inner product between $x$ and $y$. Note this is not the standard discrete Fourier transform definition over $\mathbb{Z}_N$. According to the Parseval's identity, we have

$$
  \sum_{x\in\{0,1\}^n} f(x)^2 = \sum_{y\in\{0,1\}^n} \hat{f}(y)^2  = N .
$$

```
def bitwise_dot(x, y):
    """Compute the dot product of two integers bitwise."""

    def bit_parity(i):
        n = bin(i).count("1")
        return int(n % 2)

    return bit_parity(x & y)

def fourier_transform_over_z2(v):
    """Fourier transform function over z_2^n group.

    Args:
        v: an array with 2**n elements.

    Returns:
        vs: a numpy array with same length as input.
    """
    assert len(v) & (len(v) - 1) == 0  # make sure v is 2**n long vector
    N = len(v)
    v_hat = np.array([0.0] * N)
    for y in range(N):
        for x in range(N):
            v_hat[y] += ((-1)**bitwise_dot(x, y)) * v[x]
    return v_hat / np.sqrt(N)
```

Let's have some examples in $\mathbb{Z}^2_2$. You should verify that both functions have same energy 4.

```
f = np.array([1, -1, 1, -1])
f_hat = fourier_transform_over_z2(f)
print(f"f: {list(f)} f_hat: {list(f_hat)}")

f = np.array([1, 1, 1, -1])
f_hat = fourier_transform_over_z2(f)
print(f"f: {list(f)} f_hat: {list(f_hat)}")

f = np.array([1, -1, -1, 1])
f_hat = fourier_transform_over_z2(f)
print(f"f: {list(f)} f_hat: {list(f_hat)}")
```

## Forrelation

The concept of *forrelation* is the key concept in the Fourier checking problem, which is a combination of the words "Fourier Transform" and "Correlation".
Recall the classical correlation between two vectors $u$ and $v$ is defined as
$$
  {\rm correlation}(u,v) = \frac{ \langle u, v \rangle }{\|u\| \|v\|}.
$$
Then, the forrelation between two vectors $u$ and $v$ is just the correlation between $u$ and the Fourier trasformed of $v$ -- denoted as $\widehat{v}$:
\begin{align}
  {\rm forrelation}(u, v) =& \frac{\langle u, \widehat{v} \rangle }{\|u\| \|\widehat{v}\|},\\
  =& \frac{ \langle u, \widehat{v} \rangle }{\|u\| \|v\|}.
\end{align}
where the second equality is due to the Parseval's identity.
Since in this tutorial we are interesed in Boolean function, we replace the arbitary vector $u$ and $v$ by the output of Boolean function $f$ and $g$. Now we can further simplify the above definition:
\begin{align}
  {\rm forrelation}(f, g) =& \frac{\langle f, \widehat{g} \rangle }{\|f\| \|g\|}\\
  =& \frac{1}{N} \langle f, \widehat{g}\rangle  \\
  =& \frac{1}{N} \sum_{x \in \{0,1\}^n}f(x)\widehat{g}(x)\\
  =& \frac{1}{N^{3/2}} \sum_{x, y \in \{0,1\}^n}f(x)(-1)^{x \cdot y}g(y)
\end{align}

```
def get_correlation(f, g):
    return f.dot(g) / np.linalg.norm(f) / np.linalg.norm(g)

def get_forrelation(f, g):
    g_hat = fourier_transform_over_z2(g)
    return f.dot(g_hat) / np.linalg.norm(f) / np.linalg.norm(g)
```

```
# let's see some examples to gain some insights of forrelation
f = np.array([1, -1, 1, -1])
g = np.array([1, -1, 1, -1])
print(f"Correlation: {get_correlation(f,g)}  Forrelation: {get_forrelation(f,g)}")

f = np.array([1, 1, 1, -1])
g = np.array([-1, -1, -1, 1])
print(f"Correlation: {get_correlation(f,g)}  Forrelation: {get_forrelation(f,g)}")

f = np.array([1, -1, -1, 1])
g = np.array([1, 1, 1, 1])
print(f"Correlation: {get_correlation(f,g)}  Forrelation: {get_forrelation(f,g)}")
```

# Fourier Checking Problem

Now we are prepared to discuss the *Fourier Checking* problem. Here we are given oracle access to two Boolean functions $f,g : \{0, 1\}^n \rightarrow\{-1, 1\}$. We are promised that one of the following two cases is true:

- $\langle f, g \rangle$ was drawn from the uniform distribution $\mathcal{U}$.
- $\langle f, g \rangle$ was drawn from the forrelated distribution $\mathcal{F}$ (Will be explained in more details later).

The problem is a decision problem that accepts the $\langle f, g \rangle$ if it was drawn from $\mathcal{F}$ and rejects $\langle f, g \rangle$ if it was drawn from $\mathcal{U}$.

*Note: Since $\mathcal{F}$ and $\mathcal{U}$ overlap slightly, we can only hope to succeed with overwhelming probability over the choice of  $\langle f, g \rangle$ , not for every  $\langle f, g \rangle$  pair.*

## Generate functions $f$ and $g$ from distribution $\mathcal U$ and $\mathcal F$
The forrelated distribution $\mathcal{F}$ is generated as followings.
First
choose a random real vector $v=(v_x)_{x\in\{0, 1\}^n} \in \mathbb{R}^N$, by drawing each entry independently from Gaussian distribution with zero-mean and variance 1. Then set $f(x) := {\rm sgn}(v_x)$ and $g(x) := {\rm sgn}(\widehat{v}_x)$, where $\widehat{v}_x$ the Fourier Transform of $v$ is
$$
  \widehat{v}_y := \frac{1}{\sqrt{N}} \sum_{x\in\{0,1\}^n} (-1)^{x\cdot y}v_x,
$$
and
$$
  {\rm sgn}(\alpha) := \left\{
  \begin{aligned}
    1 \;\;\; &{\rm if}\; \alpha \geq 0 \\
    -1 \;\;\;&{\rm if}\; \alpha < 0
  \end{aligned}\right.
$$
Notice, $f$ and $g$ *individually* are still uniformly random, but they are no longer independent. Now $f$ is forrelated with $g$.
For simplicity, we only consider the *PROMISE FOURIER CHECKING* problem. Under this situation, we are promised that the quatity:

\begin{align}
  p(f,g) := \left[{\rm forrelation}(f, g)\right]^2 = \frac{1}{N^3} \left(\sum_{x, y \in \{0,1\}^n}f(x)(-1)^{x \cdot y}g(y)\right)^2
\end{align}

is either at least 0.05 or at most 0.01.

```
n = 6
N = 2 ** n
```

```
# We can find a forrelated pair "as promised" through while-loop
def draw_two_distribution_from_f_set(N):
    sgn = lambda x: 1 if x >= 0 else -1
    forrelation = 0.2
    while (abs(forrelation)**2 < 0.05) and (abs(forrelation)**2 > 0.01):
        vs = np.array([np.random.normal() for _ in range(N)])
        vs_hat = fourier_transform_over_z2(vs)
        fs = np.array([sgn(v) for v in vs])
        gs = np.array([sgn(v_hat) for v_hat in vs_hat])
        forrelation = get_forrelation(fs, gs)
        correlation = get_correlation(fs, gs)
    return fs, gs, forrelation, correlation

def draw_two_distribution_from_u_set(N):
    sgn = lambda x: 1 if x >= 0 else -1
    forrelation = 0.2
    while (abs(forrelation)**2 < 0.05) and (abs(forrelation)**2 > 0.01):
        vs = np.array([np.random.normal() for _ in range(N)])
        fs = np.array([sgn(v) for v in vs])
        us = np.array([np.random.normal() for _ in range(N)])
        gs = np.array([sgn(u) for u in us])
        forrelation = get_forrelation(fs, gs)
        correlation = get_correlation(fs, gs)
    return fs, gs, forrelation, correlation
```

```
fs, gs, forrelation, correlation = draw_two_distribution_from_u_set(N)
print(f"fs: {list(fs)}")
print(f"gs: {list(gs)}")

plt.figure(figsize=(15, 5))
plt.stem(fs, use_line_collection=True)
plt.stem(gs, linefmt='--r', markerfmt='ro', use_line_collection=True)
plt.title(f"Correlation: {correlation} Forrelation: {forrelation}")
```

```
fs, gs, forrelation, correlation = draw_two_distribution_from_f_set(N)
print(f"fs: {list(fs)}")
print(f"gs: {list(gs)}")
plt.figure(figsize=(15, 5))
plt.stem(fs, use_line_collection=True)
plt.stem(gs, linefmt='--r', markerfmt="ro", use_line_collection=True)
plt.title(f"Correlation: {correlation} Forrelation: {forrelation}")
```

Typically, $\mathcal{U}$ and $\mathcal{F}$ is not obviously different from each other even after we plot the whole functions information together. However, it is not hard to show that Fourier Checking is in BQP: basically, one can prepare a uniform superposition over all $x\in\{0,1\}^n$, then query $f$, apply a quantum Fourier transform, query $g$, and
check whether one has recovered something close to the uniform superposition. On the other hand, being forrelated seems like an extremely “global” property of $f$ and $g$: one that would not be apparent from querying any small number of $f$ and $g$ values, regardless of the outcomes of those queries.

# Quantum Algorithm for Fourier Checking

Now, we present the quantum algorithm for Fourier Checking, which is quite simple actually.
It can achieve the constant error probability with $O(1)$ query complexity.

**Algorithm Description:**

First, we prepare a uniform superposition over all $x \in \{0, 1\}^n$. Then query $f$ in superposition, to create the state
$$
  \frac{1}{\sqrt{N}} \sum_{x \in \{0, 1\}^n} f(x) |x\rangle
$$
Applying Hadmard gates to all $n$ qubits, to create the state
$$
  \frac{1}{N} \sum_{x,y  \in \{0, 1\}^n} f(x) (-1)^{x\cdot y} |y\rangle
$$
Then query $g$ in superposition, to create the state
$$
  \frac{1}{N} \sum_{x,y  \in \{0, 1\}^n} f(x) (-1)^{x\cdot y}g(y) |y\rangle
$$
Then apply Hadmard gates to all $n$ qubits again, to create the state
$$
  \frac{1}{N^{3/2}} \sum_{x,y  \in \{0, 1\}^n} f(x) (-1)^{x\cdot y} g(y) (-1)^{y \cdot z}|z\rangle
$$
Finally, measure in the computational basis, and "accept" if and only if the outcome $|0\rangle^{\otimes n}$is observed.

If needed, repeat the whole algorithm $O(1)$ times to boost the success probability.

It is clear that the probablity of observation $|0\rangle^{\otimes n}$ equals the quantity:
\begin{align}
  p(f,g) := \frac{1}{N^3} \left(\sum_{x, y \in \{0,1\}^n}f(x)(-1)^{x \cdot y}g(y)\right)^2
\end{align}

It is shown in the [paper](https://arxiv.org/pdf/0910.4698.pdf) that

\begin{align}
   {\rm Pr}_{\langle f, g \rangle \sim \mathcal{U}} [p(f,g) \geq 0.01] \leq& \; \frac{100}{N} \\
   {\rm Pr}_{\langle f, g \rangle \sim \mathcal{F}} [p(f,g) \geq 0.05] \geq&\;\frac{1}{50} \\
\end{align}
This implies that the probability of the forrelation square between two functions drawing from uniform  distribution $\mathcal{U}$ having larger than 0.01  will decaying quickly when we have more number of qubits. Hence, the Promise Fourier Checking problem can be solved through simply accepting when $p(f,g) \geq 0.05$ and rejecting when $p(f,g) \leq 0.01$ with constant error probability, using $O(1)$ queries to $f$ and $g$.

## Implementation in Cirq

Above algorithm is simple and straightforward to implement. It only involves with Hadmard gates and function oracles. Implementation of oracles based on truth table in Cirq is just simple diagnal gates. To see that, let's use a oracle defined over $\mathbb{Z}_2^2$ as example:

\begin{align}
  \sum_{x\in \{0, 1\}^2} f(x)|x\rangle
  =&\frac{1}{2}\Big(f(0,0)|00\rangle + f(0,1)|01\rangle+f(1,0)|10\rangle + f(1,1) |11\rangle\Big)\\
  =&\;\;\left[ \begin{array}{cccc}
    f(0,0)   & & & \\
    &  f(1,0)  & & \\
    &   &  f(1,0)  & \\
    &  & &   f(1,1) \\
  \end{array} \right]
  \left[ \begin{array}{c}
  1/2\\
  1/2\\
  1/2\\
  1/2
  \end{array} \right]
\end{align}

It is crucial to note that the output of $f$ is either 1 or -1, so the diagonal matrix is unitary.

```
def oracle(fs, qubits):
    return cirq.MatrixGate(np.diag(fs).astype(complex))(*qubits)

def fourier_checking_algorithm(qubits, fs, gs):
    """Returns the circuit for Fourier Checking algorithm given an input."""
    yield cirq.ParallelGateOperation(cirq.H, qubits)
    yield oracle(fs, qubits)
    yield cirq.ParallelGateOperation(cirq.H, qubits)
    yield oracle(gs, qubits)
    yield cirq.ParallelGateOperation(cirq.H, qubits)
    yield cirq.measure(*qubits)

qubits = cirq.LineQubit.range(n)
fs, gs, forrelation, correlation = draw_two_distribution_from_f_set(N)
circuit = cirq.Circuit(fourier_checking_algorithm(qubits, fs, gs))
print(circuit)
```

We derived that the square forrelation between $f$ and $g$ is the same as the state of final state in circuit so we can use cirq to check it. Just remember the final state is *never* able to be obtained in reality. In simulation, it is completely doable through `final_state_vector` or the `dirac_notation` of each moment step:

```
assert np.isclose(circuit.final_state_vector()[0], forrelation)

s = cirq.Simulator()
for step in s.simulate_moment_steps(circuit):
    print(step.dirac_notation())
    print("|0> state probability to observe: ",
          np.abs(step.state_vector()[0])**2)
```

```
final_state = circuit.final_state_vector()
plt.fill_between(np.arange(len(final_state)),
                 np.abs(final_state)**2)
plt.xlabel("State of qubits")
plt.ylabel("Probability")
plt.show()
```

In reality, we can measure the state of qubits only. Each measurement will only produce one state. In order to eastimate the probability, we can do 100 repeitations and use frequency of showing 0 state as the approximation of the probability of state 0.

```
repetitions = 100
obs = s.run(circuit, repetitions=repetitions)
qubits_name = ','.join(str(i) for i in range(n))
times_zero_was_measured = len(obs.data[obs.data[qubits_name] == 0])
print(
    f"times zero state was measured from {repetitions} measurements:" +
    f"{times_zero_was_measured} - {float(times_zero_was_measured/repetitions)*100}%"
)
if float(times_zero_was_measured / repetitions) > 0.05:
    print("fs and gs is forrelated!")
```

Last, we can randomly draw the functions from either $\mathcal{U}$ or $\mathcal{F}$ set to evaluate the confusion matrix of the Fourier Checking algorithm. With the confusion matrix, you should be confident that quantum algorithm instead can solve the Fourier Checking in $O(1)$ time -- even though we need to measure 100 or 1000 times to evaluate the probability, it is irrelevant to the number of states or qubits.

```
res = pd.DataFrame()
repetitions = 100
num_rounds = 1000
for _ in range(num_rounds):
    if np.random.rand() > 0.5:
        fs, gs, _, _ = draw_two_distribution_from_f_set(N)
        source = "F set"
    else:
        fs, gs, _, _ = draw_two_distribution_from_u_set(N)
        source = "U set"

    circuit = cirq.Circuit(fourier_checking_algorithm(qubits, fs, gs))
    obs = s.run(circuit, repetitions=repetitions)
    times_zero_was_measured = len(obs.data[obs.data[','.join(
        str(i) for i in range(n))] == 0])
    if times_zero_was_measured / repetitions > 0.05:
        res = res.append({
            "Source": source,
            "Decision": "accept",
            "Count": 1
        }, ignore_index=True)
    else:
        res = res.append({
            "Source": source,
            "Decision": "reject",
            "Count": 1
        }, ignore_index=True)
confusion = res.pivot_table(index="Source", columns="Decision", values="Count", aggfunc="sum")
# Translate the counts into percentage
confusion.div(confusion.sum(axis=1), axis=0).apply(lambda x: round(x, 4) * 100)
```

# Query Complexity of Quantum Algorithm Versus the Classical one

After we present the quantum algorithm, it is time to look back to classical algorithm for comparison. It is not hard to give a classical algorithm that solves Fourier Checking using $O(\sqrt{N}) = O(2^{n/2})$ queries, which is similar as we did in the BPP section.

For some $K=\Theta(\sqrt{N})$, first choose sets $X=\{x_1,\ldots, x_K\}$ and $Y = \{y_1, \ldots, y_K\}$ of $n-$bit strings uniformly at random. Then query $f(x_i)$ and $g(y_i)$ for all $i \in [K]$. Finally, compute
$$
  Z := \sum_{i,j=1}^K f(x_i) (-1)^{x_i\cdot y_j} g(y_j)
$$
accept if $|Z|$ is greater than some cutoff $cK$, and reject otherwise. For suitable $K$ and $c$, one can show that this algorithm accepts a forrelated $\langle f, g \rangle$ pair with probability at least $2/3$. Comparing this with $O(1)$ complexity in quantum query, we can see an exponential speedup by using quantum computer.

# Further reading

The Forrelation problem was originally introduced in [1]. Later, a **$k$-fold Forrelation** problem was introduced in [2], which considered the forrelation between $k$ oracle functions. In that paper, it also improved the proof about the seperation of forrelation problem from $\Omega\left(\sqrt[4]{N}\right)$ to $\tilde{\Omega}\left(\sqrt{N}\right)$, which can not be further improved. Thus, resolving an open question of Buhrman et al., there is NO partial
Boolean function whose quantum query complexity is constant and whose randomized query complexity is linear.

[1] *Scott Aaronson. BQP and the Polynomial Hierarchy. STOC ’10, page 141–150, New York, NY, USA, 2010.* [arXiv](https://arxiv.org/pdf/0910.4698.pdf)

[2] *Scott Aaronson and Andris Ambainis. Forrelation: A problem that optimally separates quantum
from classical computing. SIAM J. Comput., 47(3):982–1038, 2018.* [arXiv]( https://arxiv.org/pdf/1411.5729.pdf)


```python
%matplotlib inline
```

Advanced: Making Dynamic Decisions and the Bi-LSTM CRF
======================================================

Dynamic versus Static Deep Learning Toolkits
--------------------------------------------

Pytorch is a *dynamic* neural network kit. Another example of a dynamic
kit is `Dynet <https://github.com/clab/dynet>`__ (I mention this because
working with Pytorch and Dynet is similar. If you see an example in
Dynet, it will probably help you implement it in Pytorch). The opposite
is the *static* tool kit, which includes Theano, Keras, TensorFlow, etc.
The core difference is the following:

* In a static toolkit, you define
  a computation graph once, compile it, and then stream instances to it.
* In a dynamic toolkit, you define a computation graph *for each
  instance*. It is never compiled and is executed on-the-fly

Without a lot of experience, it is difficult to appreciate the
difference. One example is to suppose we want to build a deep
constituent parser. Suppose our model involves roughly the following
steps:

* We build the tree bottom up
* Tag the root nodes (the words of the sentence)
* From there, use a neural network and the embeddings
  of the words to find combinations that form constituents. Whenever you
  form a new constituent, use some sort of technique to get an embedding
  of the constituent. In this case, our network architecture will depend
  completely on the input sentence. In the sentence "The green cat
  scratched the wall", at some point in the model, we will want to combine
  the span $(i,j,r) = (1, 3, \text{NP})$ (that is, an NP constituent
  spans word 1 to word 3, in this case "The green cat").

However, another sentence might be "Somewhere, the big fat cat scratched
the wall". In this sentence, we will want to form the constituent
$(2, 4, NP)$ at some point. The constituents we will want to form
will depend on the instance. If we just compile the computation graph
once, as in a static toolkit, it will be exceptionally difficult or
impossible to program this logic. In a dynamic toolkit though, there
isn't just 1 pre-defined computation graph. There can be a new
computation graph for each instance, so this problem goes away.

Dynamic toolkits also have the advantage of being easier to debug and
the code more closely resembling the host language (by that I mean that
Pytorch and Dynet look more like actual Python code than Keras or
Theano).

Bi-LSTM Conditional Random Field Discussion
-------------------------------------------

For this section, we will see a full, complicated example of a Bi-LSTM
Conditional Random Field for named-entity recognition. The LSTM tagger
above is typically sufficient for part-of-speech tagging, but a sequence
model like the CRF is really essential for strong performance on NER.
Familiarity with CRF's is assumed. Although this name sounds scary, all
the model is is a CRF but where an LSTM provides the features. This is
an advanced model though, far more complicated than any earlier model in
this tutorial. If you want to skip it, that is fine. To see if you're
ready, see if you can:

-  Write the recurrence for the viterbi variable at step i for tag k.
-  Modify the above recurrence to compute the forward variables instead.
-  Modify again the above recurrence to compute the forward variables in
   log-space (hint: log-sum-exp)

If you can do those three things, you should be able to understand the
code below. Recall that the CRF computes a conditional probability. Let
$y$ be a tag sequence and $x$ an input sequence of words.
Then we compute

\begin{align}P(y|x) = \frac{\exp{(\text{Score}(x, y)})}{\sum_{y'} \exp{(\text{Score}(x, y')})}\end{align}

Where the score is determined by defining some log potentials
$\log \psi_i(x,y)$ such that

\begin{align}\text{Score}(x,y) = \sum_i \log \psi_i(x,y)\end{align}

To make the partition function tractable, the potentials must look only
at local features.

In the Bi-LSTM CRF, we define two kinds of potentials: emission and
transition. The emission potential for the word at index $i$ comes
from the hidden state of the Bi-LSTM at timestep $i$. The
transition scores are stored in a $|T|x|T|$ matrix
$\textbf{P}$, where $T$ is the tag set. In my
implementation, $\textbf{P}_{j,k}$ is the score of transitioning
to tag $j$ from tag $k$. So:

\begin{align}\text{Score}(x,y) = \sum_i \log \psi_\text{EMIT}(y_i \rightarrow x_i) + \log \psi_\text{TRANS}(y_{i-1} \rightarrow y_i)\end{align}

\begin{align}= \sum_i h_i[y_i] + \textbf{P}_{y_i, y_{i-1}}\end{align}

where in this second expression, we think of the tags as being assigned
unique non-negative indices.

If the above discussion was too brief, you can check out
`this <http://www.cs.columbia.edu/%7Emcollins/crf.pdf>`__ write up from
Michael Collins on CRFs.

Implementation Notes
--------------------

The example below implements the forward algorithm in log space to
compute the partition function, and the viterbi algorithm to decode.
Backpropagation will compute the gradients automatically for us. We
don't have to do anything by hand.

The implementation is not optimized. If you understand what is going on,
you'll probably quickly see that iterating over the next tag in the
forward algorithm could probably be done in one big operation. I wanted
to code to be more readable. If you want to make the relevant change,
you could probably use this tagger for real tasks.

```python
# Author: Robert Guthrie

import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(1)
```

Helper functions to make the code more readable.

```python
def argmax(vec):
    # return the argmax as a python int
    _, idx = torch.max(vec, 1)
    return idx.item()

def prepare_sequence(seq, to_ix):
    idxs = [to_ix[w] for w in seq]
    return torch.tensor(idxs, dtype=torch.long)

# Compute log sum exp in a numerically stable way for the forward algorithm
def log_sum_exp(vec):
    max_score = vec[0, argmax(vec)]
    max_score_broadcast = max_score.view(1, -1).expand(1, vec.size()[1])
    return max_score + \
        torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))
```

Create model

```python
class BiLSTM_CRF(nn.Module):

    def __init__(self, vocab_size, tag_to_ix, embedding_dim, hidden_dim):
        super(BiLSTM_CRF, self).__init__()
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        self.tag_to_ix = tag_to_ix
        self.tagset_size = len(tag_to_ix)

        self.word_embeds = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2,
                            num_layers=1, bidirectional=True)

        # Maps the output of the LSTM into tag space.
        self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)

        # Matrix of transition parameters.  Entry i,j is the score of
        # transitioning *to* i *from* j.
        self.transitions = nn.Parameter(
            torch.randn(self.tagset_size, self.tagset_size))

        # These two statements enforce the constraint that we never transfer
        # to the start tag and we never transfer from the stop tag
        self.transitions.data[tag_to_ix[START_TAG], :] = -10000
        self.transitions.data[:, tag_to_ix[STOP_TAG]] = -10000

        self.hidden = self.init_hidden()

    def init_hidden(self):
        return (torch.randn(2, 1, self.hidden_dim // 2),
                torch.randn(2, 1, self.hidden_dim // 2))

    def _forward_alg(self, feats):
        # Do the forward algorithm to compute the partition function
        init_alphas = torch.full((1, self.tagset_size), -10000.)
        # START_TAG has all of the score.
        init_alphas[0][self.tag_to_ix[START_TAG]] = 0.

        # Wrap in a variable so that we will get automatic backprop
        forward_var = init_alphas

        # Iterate through the sentence
        for feat in feats:
            alphas_t = []  # The forward tensors at this timestep
            for next_tag in range(self.tagset_size):
                # broadcast the emission score: it is the same regardless of
                # the previous tag
                emit_score = feat[next_tag].view(
                    1, -1).expand(1, self.tagset_size)
                # the ith entry of trans_score is the score of transitioning to
                # next_tag from i
                trans_score = self.transitions[next_tag].view(1, -1)
                # The ith entry of next_tag_var is the value for the
                # edge (i -> next_tag) before we do log-sum-exp
                next_tag_var = forward_var + trans_score + emit_score
                # The forward variable for this tag is log-sum-exp of all the
                # scores.
                alphas_t.append(log_sum_exp(next_tag_var).view(1))
            forward_var = torch.cat(alphas_t).view(1, -1)
        terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]
        alpha = log_sum_exp(terminal_var)
        return alpha

    def _get_lstm_features(self, sentence):
        self.hidden = self.init_hidden()
        embeds = self.word_embeds(sentence).view(len(sentence), 1, -1)
        lstm_out, self.hidden = self.lstm(embeds, self.hidden)
        lstm_out = lstm_out.view(len(sentence), self.hidden_dim)
        lstm_feats = self.hidden2tag(lstm_out)
        return lstm_feats

    def _score_sentence(self, feats, tags):
        # Gives the score of a provided tag sequence
        score = torch.zeros(1)
        tags = torch.cat([torch.tensor([self.tag_to_ix[START_TAG]], dtype=torch.long), tags])
        for i, feat in enumerate(feats):
            score = score + \
                self.transitions[tags[i + 1], tags[i]] + feat[tags[i + 1]]
        score = score + self.transitions[self.tag_to_ix[STOP_TAG], tags[-1]]
        return score

    def _viterbi_decode(self, feats):
        backpointers = []

        # Initialize the viterbi variables in log space
        init_vvars = torch.full((1, self.tagset_size), -10000.)
        init_vvars[0][self.tag_to_ix[START_TAG]] = 0

        # forward_var at step i holds the viterbi variables for step i-1
        forward_var = init_vvars
        for feat in feats:
            bptrs_t = []  # holds the backpointers for this step
            viterbivars_t = []  # holds the viterbi variables for this step

            for next_tag in range(self.tagset_size):
                # next_tag_var[i] holds the viterbi variable for tag i at the
                # previous step, plus the score of transitioning
                # from tag i to next_tag.
                # We don't include the emission scores here because the max
                # does not depend on them (we add them in below)
                next_tag_var = forward_var + self.transitions[next_tag]
                best_tag_id = argmax(next_tag_var)
                bptrs_t.append(best_tag_id)
                viterbivars_t.append(next_tag_var[0][best_tag_id].view(1))
            # Now add in the emission scores, and assign forward_var to the set
            # of viterbi variables we just computed
            forward_var = (torch.cat(viterbivars_t) + feat).view(1, -1)
            backpointers.append(bptrs_t)

        # Transition to STOP_TAG
        terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]
        best_tag_id = argmax(terminal_var)
        path_score = terminal_var[0][best_tag_id]

        # Follow the back pointers to decode the best path.
        best_path = [best_tag_id]
        for bptrs_t in reversed(backpointers):
            best_tag_id = bptrs_t[best_tag_id]
            best_path.append(best_tag_id)
        # Pop off the start tag (we dont want to return that to the caller)
        start = best_path.pop()
        assert start == self.tag_to_ix[START_TAG]  # Sanity check
        best_path.reverse()
        return path_score, best_path

    def neg_log_likelihood(self, sentence, tags):
        feats = self._get_lstm_features(sentence)
        forward_score = self._forward_alg(feats)
        gold_score = self._score_sentence(feats, tags)
        return forward_score - gold_score

    def forward(self, sentence):  # dont confuse this with _forward_alg above.
        # Get the emission scores from the BiLSTM
        lstm_feats = self._get_lstm_features(sentence)

        # Find the best path, given the features.
        score, tag_seq = self._viterbi_decode(lstm_feats)
        return score, tag_seq
```

Run training

```python
START_TAG = "<START>"
STOP_TAG = "<STOP>"
EMBEDDING_DIM = 5
HIDDEN_DIM = 4

# Make up some training data
training_data = [(
    "the wall street journal reported today that apple corporation made money".split(),
    "B I I I O O O B I O O".split()
), (
    "georgia tech is a university in georgia".split(),
    "B I O O O O B".split()
)]

word_to_ix = {}
for sentence, tags in training_data:
    for word in sentence:
        if word not in word_to_ix:
            word_to_ix[word] = len(word_to_ix)

tag_to_ix = {"B": 0, "I": 1, "O": 2, START_TAG: 3, STOP_TAG: 4}

model = BiLSTM_CRF(len(word_to_ix), tag_to_ix, EMBEDDING_DIM, HIDDEN_DIM)
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=1e-4)

# Check predictions before training
with torch.no_grad():
    precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
    precheck_tags = torch.tensor([tag_to_ix[t] for t in training_data[0][1]], dtype=torch.long)
    print(model(precheck_sent))

# Make sure prepare_sequence from earlier in the LSTM section is loaded
for epoch in range(
        300):  # again, normally you would NOT do 300 epochs, it is toy data
    for sentence, tags in training_data:
        # Step 1. Remember that Pytorch accumulates gradients.
        # We need to clear them out before each instance
        model.zero_grad()

        # Step 2. Get our inputs ready for the network, that is,
        # turn them into Tensors of word indices.
        sentence_in = prepare_sequence(sentence, word_to_ix)
        targets = torch.tensor([tag_to_ix[t] for t in tags], dtype=torch.long)

        # Step 3. Run our forward pass.
        loss = model.neg_log_likelihood(sentence_in, targets)

        # Step 4. Compute the loss, gradients, and update the parameters by
        # calling optimizer.step()
        loss.backward()
        optimizer.step()

# Check predictions after training
with torch.no_grad():
    precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
    print(model(precheck_sent))
# We got it!
```

Exercise: A new loss function for discriminative tagging
--------------------------------------------------------

It wasn't really necessary for us to create a computation graph when
doing decoding, since we do not backpropagate from the viterbi path
score. Since we have it anyway, try training the tagger where the loss
function is the difference between the Viterbi path score and the score
of the gold-standard path. It should be clear that this function is
non-negative and 0 when the predicted tag sequence is the correct tag
sequence. This is essentially *structured perceptron*.

This modification should be short, since Viterbi and score\_sentence are
already implemented. This is an example of the shape of the computation
graph *depending on the training instance*. Although I haven't tried
implementing this in a static toolkit, I imagine that it is possible but
much less straightforward.

Pick up some real data and do a comparison!