How to perform a range on a Theano's TensorVariable?

Example:

import theano.tensor as T
from theano import function

constant = T.dscalar('constant')
n_iters = T.dscalar('n_iters')
start = T.dscalar('start')
result = start  

for iter in range(n_iters):
    result = start + constant

f = function([start, constant, n_iters], result)
print('f(0,2,5): {0}'.format(f(1,2)))

returns the error:

Traceback (most recent call last):
  File "test_theano.py", line 9, in <module>
    for iter in range(n_iters):
TypeError: range() integer end argument expected, got TensorVariable.

What is the correct way to use a range on a Theano's TensorVariable?

Solution

It's not clear what this code is attempting to do because even if the loop worked, it wouldn't compute anything useful: result will always equal the sum of start and constant irrespective of the number of iterations.

I'll assume that the intention was to compute result like this:

for iter in range(n_iters):
    result = result + constant

The problem is that you're mixing symbolic, delayed execution, Theano operations with non-symbolic, immediately executed, Python operations.

range is a Python function that expects a Python integer parameter but you're providing a Theano symbolic value (n_iters, which has a double type instead of an integer type but let's assume it's actually an iscalar instead of a dscalar). As far as Python is concerned all Theano symbolic tensors are just objects: instances of a class type somewhere within the Theano library; they are most assuredly not integers. Even if you squint and try to pretend a Theano iscalar looks like a Python integer, it still doesn't work because Python operations execute immediately which means n_iters needs to have a value immediately. Theano on the other hand doesn't have a value for any iscalar until one is provided by calling a compiled Theano function (or via eval).

To create a symbolic range, you can use theano.tensor.arange which operates just like NumPy's arange, but symbolically.

Example:

import theano.tensor as T
from theano import function

my_range_max = T.iscalar('my_range_max')
my_range = T.arange(my_range_max)

f = function([my_range_max], my_range)
print('f(10): {0}'.format(f(10)))

outputs:

f(10): [0 1 2 3 4 5 6 7 8 9]

By making n_iters a symbolic variable you are implicitly saying "I don't know how many iterations there need to be in this loop until a value is provided for n_iters later". That being the case, you must use a symbolic loop instead of a Python for loop. In the latter case you must tell Python how many times to iterate now, you can't delay that decision until a value is provided for n_iters later. To solve this you need to switch to a symbolic loop, which is provided by Theano's scan operator.

Here's the code changed to use scan (as well as the other assumed changes).

import theano
import theano.tensor as T
from theano import function

constant = T.dscalar('constant')
n_iters = T.iscalar('n_iters')
start = T.dscalar('start')

results, _ = theano.scan(lambda result, constant: result + constant,
                         outputs_info=[start], non_sequences=[constant], n_steps=n_iters)

f = function([start, constant, n_iters], results[-1])
print('f(0,2,5): {0}'.format(f(0, 2, 5)))