Control flow in Tensorflow

In this tutorial, we're going to use control flow operations, i.e. if statements and while loops. So before we tackle the bigger problem, let's learn how these control flow operations are implemented in Tensorflow.


In [ ]:
import tensorflow as tf
import numpy as np
import pandas as pd
%matplotlib inline

if statements in Tensorflow: tf.cond

Let's start Tensorflow's 'if' statement: tf.cond ('cond' stands for 'conditional execution')


In [ ]:
# clear previous computation graph
tf.reset_default_graph()

# some dummy data
x = tf.constant(7)
y = tf.constant(13)


max_xy = tf.cond(x >= y, lambda: x, lambda: y)


with tf.Session() as s:
    print s.run(max_xy)

Note: Also have a look at tf.case for more general switch statements.

ExerciseWrite a function using tf.cond that takes a scalar and returns its modulus (without using tf.abs).


In [ ]:
# clear previous computation graph
tf.reset_default_graph()

# some test cases
x = tf.constant(...)
y = tf.constant(...)


def my_abs(x):
    '<your code here>'


with tf.Session() as s:
    '<your code here>'

In [ ]:
# %load sol/ex_cond_abs.py

Traversing sequences in Tensorflow: tf.map_fn and tf.while_loop

Let's consider an array $x$, e.g.

x = [0, 1, 2, 3]

There are essentially two distinct ways of traversing an array.

Type 1: map operations

The input for each iteration $i$ depends only on the component $x_i$. This is often called a map operation, which are trivially parallelizable (think of the map-reduce paradigm). In Tensorflow, this is implemented in tf.map_fn.

Type 2: loops

Here, the input for each consecutive iteration $i$ depends on the values computed in the previous iteration $i-1$. In other words, there is a definite notion of order of the sequence of operations. This kind of operation is implemented in Tensorflow in tf.while_loop.

In the following two exercises we're going to learn how to use these two operations.

Exercise#### Use tf.map_fn to broadcasting unequal-sized tensors.

Now suppose we have the following two tensors:

a = tf.ones([2, 3, 5])
b = tf.ones([7, 5])

and we would like compute the point-wise difference along the last axes (of length 5) in such a way that we get a rank-4 tensor of shape [7, 2, 3, 5]. In terms of components, this would be:

c[i,j,k,l] = a[j,k,l] - b[i,l]

Now, the problem arises when we naively try:

c = a - b

which raises an error saying:

Dimensions must be equal, but are 3 and 7 for 'sub' (op: 'Sub') with input shapes: [2,3,5], [7,5].

Now we do know that the broadcast works fine if we have a rank-1 tensor (vector) instead of a rank-2 tensor (matrix). In other words, if we select only one row of b, the broadcasting will work just fine, i.e.

a - b[0]  # this works!

Using this information, come up with a way to calculate $c$ with tf.map_fn.

Note: The output tensor c must have shape (7, 2, 3, 5).


In [ ]:
tf.reset_default_graph()


a = tf.ones([2, 3, 5])
b = tf.ones([7, 5])

    
c = '<your code here>'

In [ ]:
# %load sol/ex_map_fn_broadcasting.py

Exercise

In this exercise, we're going to use tf.while_loop to compute the sum and the cumulative sum of an array of numbers:

x = tf.constant([7, 0, 42, 1, 13, 4])

A) A simple while loop: sum.

Write a while loop that computes the sum of the entries of x.


In [ ]:
tf.reset_default_graph()

# input
x_len = 6
x = tf.constant([7, 0, 42, 1, 13, 4])


def cond(i, acc):
    '<your code here>'


def body(i, acc):
    '<your code here>'


# initial values for the loop variables
loop_vars = '<your code here>'

# compute loop
'<your code here>'


with tf.Session() as s:
    '<your code here>'

In [ ]:
# %load sol/ex_while_loop_sum.py

Now that we have the simple stuff out of the way, we're ready to tackle the slightly more involved case in which the accumulator isn't a scalar.

B) A slightly trickier while loop: cumsum.

In this part of the exercise, we'll build on our previous solution to calculate the cumulative sum ox x.

The main idea is the following. You start with an empty vector:

y = tf.constant([], dtype=...)

Then, as you're computing the sum the way you did in part A, you keep appending each new entry to y with every iteration. To do the 'appending', you could use e.g. tf.concat in the following way:

a = tf.constant([7, 1, 13])  # vector
b = tf.constant(11)          # scalar

# append scalar b to vector a
a = tf.concat([a, tf.expand_dims(b, axis=0)], axis=0)

Using these instructions, calculate the cumulative sum using tf.while_loop.


In [ ]:
tf.reset_default_graph()

# input
x_len = 6
x = tf.constant([7, 0, 42, 1, 13, 4])


def cond(i, acc, y):
    '<your code here>'


def body(i, acc, y):
    '<your code here>'


# initial values for the loop variables
loop_vars = '<your code here>'

# specify dynamic shape invariant for y
shape_invariants = '<your code here>'

# compute the loop
'<your code here>'


with tf.Session() as s:
    '<your code here>'

In [ ]:
# %load sol/ex_while_loop_cumsum.py