Search code examples
theanopymc3theano.scan

Using theano.scan within PyMC3 gives TypeError: slice indices must be integers or None or have an __index__ method

I would like to to use theano.scan within pymc3. I run into problems when I add more than two variables as sequences. Here is a simple example:

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

a = np.ones(5)
b = np.ones(5)

basic_model = pm.Model()
with basic_model:
    a_plus_b, _ = theano.scan(fn=lambda a, b: a + b, sequences=[a, b])

results in the following error:

Traceback (most recent call last):
File "StackOverflowExample.py", line 23, in <module>
sequences=[a, b])
File "\Anaconda3\lib\site-packages\theano\scan_module\scan.py", line 586, in scan
scan_seqs = [seq[:actual_n_steps] for seq in scan_seqs]
File "\Anaconda3\lib\site-packages\theano\scan_module\scan.py", line 586, in <listcomp>
scan_seqs = [seq[:actual_n_steps] for seq in scan_seqs]
TypeError: slice indices must be integers or None or have an __index__ method

However, when I run the same theano.scan outside a pymc model block, everything works fine:

a = T.vector('a')
b = T.vector('b')
a_plus_b, update = theano.scan(fn=lambda a, b: a + b, sequences=[a, b])
a_plus_b_function = theano.function(inputs=[a, b], outputs=a_plus_b, updates=update)

a = np.ones(5)
b = np.ones(5)
print(a_plus_b_function(a, b))

prints [2. 2. 2. 2. 2.], like it should.

In addition, the problem seems to be specific to adding more than one sequences. Everything works just fine when there is one variable in sequences and one in non-sequences. The following code works:

a = np.ones(5)
c = 2

basic_model = pm.Model()
with basic_model:

    a_plus_c, _ = theano.scan(fn=lambda a, c: a + c, sequences=[a], non_sequences=[c])
    a_plus_c_print = T.printing.Print('a_plus_c')(a_plus_c)

prints a_plus_c __str__ = [ 3. 3. 3. 3. 3.], as expected.

Note: I can't just use a + b instead of theano.scan because my actual function is more complex. I actually want to have something like this:

rewards = np.array([1, 1, 1, 1])  # reward (1) or no reward (0)
choices = np.array([1, 0, 1, 0])  # action left (1) or right (0)
Q_old = 0  # initial Q-value
alpha = 0.1  # learning rate

def update_Q(reward, choice, Q_old, alpha):
    return Q_old + choice * alpha * (reward - Q_old)

Q_left, _ = theano.scan(fn=update_Q,
                        sequences=[rewards, choices],
                        outputs_info=[Q_old],
                        non_sequences=[alpha])
Solution

  • Turns out it was a simple mistake! Everything is working as soon as I define a and b as tensor variables. Adding those two lines did the job:

    a = T.as_tensor_variable(np.ones(5))
b = T.as_tensor_variable(np.ones(5))