glS
glS

Reputation: 1299

How to scan through all the elements of a matrix with theano?

TL;DR:

What is the theano.scan equivalent of:

M = np.arange(9).reshape(3, 3)
for i in range(M.shape[0]):
    for j in range(M.shape[1]):
        M[i, j] += 5
M

possibly (if doable) without using nested scans?

Note that this question does not want to be specifically about how to apply an operation elementwise to a matrix, but more generally on how to implement with theano.scan a nested looping construct like the above.


Long version:

theano.scan (or equivalently in this case, theano.map) allows to map a function looping through multiple indices, by simply providing a sequence of elements to the sequences arguments, with something like

import theano
import theano.tensor as T
M = T.dmatrix('M')
def map_func(i, j, matrix):
    return matrix[i, j] + i * j
results, updates = theano.scan(map_func,
            sequences=[T.arange(M.shape[0]), T.arange(M.shape[1])],
            non_sequences=[M])
f = theano.function(inputs=[M], outputs=results)
f(np.arange(9).reshape(3, 3))
# 

which is roughly equivalent to a python loop of the form:

M = np.arange(9).reshape(3, 3)
for i, j in zip(np.arange(M.shape[0]), np.arange(M.shape[1])):
    M[i, j] += 5
M

which increases by 5 all the elements in the diagonal of M.

But what if I want to find the theano.scan equivalent of:

M = np.arange(9).reshape(3, 3)
for i in range(M.shape[0]):
    for j in range(M.shape[1]):
        M[i, j] += 5
M

possibly without nesting scan?

One way is of course to flatten the matrix, scan through the flattened elements, and then reshape it to the original shape, with something like

import theano
import theano.tensor as T
M = T.dmatrix('M')
def map_func(i, X):
    return X[i] + .5
M_flat = T.flatten(M)
results, updates = theano.map(map_func,
                              sequences=T.arange(M.shape[0] * M.shape[1]),
                              non_sequences=M_flat)
final_M = T.reshape(results, M.shape)
f = theano.function([M], final_M)
f([[1, 2], [3, 4]])

but is there a better way that doesn't involve explicitly flattening and reshaping the matrix?

Upvotes: 2

Views: 908

Answers (1)

glS
glS

Reputation: 1299

Here is an example on how this kind of thing can be achieve using nested theano.scan calls. In this example we add the number 3.141 to every element of a matrix, effectively simulating in a convoluted way the output of H + 3.141:

H = T.dmatrix('H')
def fn2(col, row, matrix):
    return matrix[row, col] + 3.141

def fn(row, matrix):
    res, updates = theano.scan(fn=fn2,
                               sequences=T.arange(matrix.shape[1]),
                               non_sequences=[row, matrix])
    return res

results, updates = theano.scan(fn=fn,
                               sequences=T.arange(H.shape[0]),
                               non_sequences=[H])
f = theano.function([H], results)
f([[0, 1], [2, 3]])
# array([[ 3.141,  4.141],
#        [ 5.141,  6.141]])

As another example, let us add to each element of a matrix the product of its row and column indices:

H = T.dmatrix('H')
def fn2(col, row, matrix):
    return matrix[row, col] + row * col

def fn(row, matrix):
    res, updates = theano.scan(fn=fn2,
                               sequences=T.arange(matrix.shape[1]),
                               non_sequences=[row, matrix])
    return res

results, updates = theano.scan(fn=fn,
                               sequences=T.arange(H.shape[0]),
                               non_sequences=[H])
f = theano.function([H], results)
f(np.arange(9).reshape(3, 3))
# Out[2]:array([[  0.,   1.,   2.],
#               [  3.,   5.,   7.],
#               [  6.,   9.,  12.]])

Upvotes: 1

Related Questions