Theano stack matrices programmatically?

Question

I have the following code which stacks 2 matrices into a 3D tensor.

import theano
import theano.tensor as T
A = T.matrix("A")
B = theano.tensor.stack(A, A)
f = theano.function(inputs=[A], outputs=B)
print f([range(10)]*2)

However, I do not know how many times I need to stack the matrix in advance. For example the fourth line of code may be:

B = theano.tensor.stack(A, A, A)
B = theano.tensor.stack(A, A, A, A)
etc...

Is there a theano function to duplicate a matrix n times:

theano.some_function(A, 3) = theano.tensor.stack(A, A, A)

Then I can pass that 3, as an argument to the theano function f. Is this possible? I looked into broadcasting but broadcasting does not explicitly change dimensionality/stack.

Answer 1

After digging long and hard through the theano documentation I have found the solution:

import theano
import theano.tensor as T
A = T.matrix("A")
B = [A]
C = theano.tensor.extra_ops.repeat(B, 3, axis=0)
f = theano.function(inputs=[A], outputs=C)
print f([range(10)]*2)

is equivalent to:

import theano
import theano.tensor as T
A = T.matrix("A")
B = theano.tensor.stack(A, A, A)
f = theano.function(inputs=[A], outputs=B)
print f([range(10)]*2)

except we can now choose the number of repeats programatically as the second argument to: theano.tensor.extra_ops.repeat

Answer 2

Here is an example using broadcasting

import theano
import theano.tensor as T
import numpy as np

A = T.fmatrix()
n = T.iscalar()

ones = T.ones((n, 1, 1))

stackedA = ones * A[np.newaxis, :, :]

f = theano.function([A, n], stackedA)

a = np.arange(30).reshape(5, 6).astype('float32')
nn = 3

r = f(a, nn)

print r.shape  # outputs (3, 4, 5)
print (r == a[np.newaxis]).all()  # outputs True

This approach can help the compiler avoid tiling if it can optimize that away.