Constructing a matrix within each step of `pytensor.scan`

1

I’d like to use scan to iteratively multiply a matrix onto a vector. At each step, the matrix elements are first constructed from the values of the scan routine’s sequences and non_sequences, and then the matrix is multiplied onto the vector. The state of the vector is updated from one iteration to the next via outputs_info (scan tutorial e.g. 4).

An attempt at the code is pasted below. This attempt yields TypeError: Unsupported dtype for TensorType: object, thrown at the call to pt.tensor.as_tensor_variable().

def step(t, 
         tau_01, tau_10, a_01, a_10,
         p):
     """
     step(sequences,
          non_sequences,
          outputs_info)
     """
     # Transition rates in effect at time `t` within `times`
     gamma_01 = a_01 * pt.tensor.math.exp(-t/tau_01)
     gamma_10 = a_10 * pt.tensor.math.exp(-t/tau_10)

     # Transition matrix
     T = pt.tensor.as_tensor_variable(
          np.array([[-gamma_01,  gamma_10],
                    [ gamma_01, -gamma_10]])
     )
     
     return pt.tensor.dot(T,p)

# sequences
times = pt.tensor.dvector('times')
# non_sequences
tau_01 = pt.tensor.dscalar('tau_01')
tau_10 = pt.tensor.dscalar('tau_10')
a_01 = pt.tensor.dscalar('a_01')
a_10 = pt.tensor.dscalar('a_10')
# outputs_info
p = pt.tensor.dvector('p')

output, updates = pt.scan(fn=step, 
                          sequences=[times],
                          non_sequences=[tau_01, tau_10, a_01, a_10],
                          outputs_info=[p])

f = pt.function(inputs=[times, 
                        tau_01, tau_10, a_01, a_10,
                        p], 
                outputs=output,
                updates=updates)

# sequences
times_value = np.arange(0, 100+1, 1)
# non_sequences
tau_01_value = 40
tau_10_value = 20
a_01_value = 5
a_10_value = 50
# outputs_info
p_value = np.asarray([1, 0]) # [state-0, state-1]

result = f(times_value, 
           tau_01_value, tau_10_value, a_01_value, a_10,
           p_value)

This question is similar to Learning pytensor scan function utilizing a simple population model in the sense that I’m attempting to do something like forward simulation of a population model, but with the added complication of computing new elements of the transition matrix at each step.

What might I have wrong here? Is there a recommended way of constructing a matrix within the step function of a scan routine?

1 Like
2

Don’t use numpy functions (e.g. np.array) inside pytensor code. Everything inside the inner function is symbolic, so you don’t need any as_tensor_variable, just do the operations you want to do. You can use e.g. pt.stacklists:

import pytensor.tensor as pt
a, b, c, d = pt.dscalars(*list('abcd'))
x = pt.stacklists([[a, b], [c, d]])
x.eval({a:0, b:1, c:2, d:3})
# array([[0., 1.],
#       [2., 3.]])

You also don’t need a scan for this particular example, you could just use broadcasting. I assume its a toy example, but always keep this in mind, because vectorization is way faster than looping.

1 Like
3

@jessegrabowski, Thanks for your reply.

Using your snippet, I can compute a single step of my desired routine (see code block below). Note that t is a scalar entry from the times vector.

# sequences
times = pt.tensor.dvector('times')
t = pt.tensor.dscalar('t') # entry from `times` sequence
# non_sequences
dt = pt.tensor.dscalar('dt')
tau_01 = pt.tensor.dscalar('tau_01')
tau_10 = pt.tensor.dscalar('tau_10')
a_01 = pt.tensor.dscalar('a_01')
a_10 = pt.tensor.dscalar('a_10')
# outputs_info
p = pt.tensor.dvector('p')

# sequences
times_value = np.arange(0, 100+1, 1)
t_value = times_value[0] # first entry from `times` sequence
# non_sequences
dt_value = 1
tau_01_value = 40
tau_10_value = 20
a_01_value = 5
a_10_value = 50
# outputs_info
p_value = np.array([1, 0]) # [state-0, state-1]

# Computation of single step
gamma_01 = a_01 * pt.tensor.math.exp(-t/tau_01)
gamma_10 = a_10 * pt.tensor.math.exp(-t/tau_10)
T = pt.tensor.stacklists(
          [[-gamma_01,  gamma_10],
           [ gamma_01, -gamma_10]]
     ) * dt
p_new = pt.tensor.dot(T,p) + p

# Evaluate
p_new.eval({t:t_value, 
            dt:dt_value, tau_01:tau_01_value, tau_10:tau_10_value, a_01:a_01_value, a_10:a_10_value,
            p:p_value})


>> array([-4.,  5.])

However, when I take that working code and place it within the step function of scan, I get the error

ValueError: When compiling the inner function of scan (the function called by scan in each of its iterations) the following error has been encountered: The initial state (outputs_info in scan nomenclature) of variable SetSubtensor{:stop}.0 (argument number 1) has 2 dimension(s), while the corresponding variable in the result of the inner function of scan (fn) has 2 dimension(s) (it should be one less than the initial state). ...

Here is the errant attempt at the scan routine:

def step(t, 
         dt, tau_01, tau_10, a_01, a_10,
         p):
     """
     step(sequences,
          non_sequences,
          outputs_info)
     """
     # Computation of single step
     gamma_01 = a_01 * pt.tensor.math.exp(-t/tau_01)
     gamma_10 = a_10 * pt.tensor.math.exp(-t/tau_10)
     T = pt.tensor.stacklists(
               [[-gamma_01,  gamma_10],
               [ gamma_01, -gamma_10]]
          ) * dt
     p_new = pt.tensor.dot(T,p) + p

     return p_new

# sequences
times = pt.tensor.dvector('times')
t = pt.tensor.dscalar('t') # entry from `times` sequence
# non_sequences
dt = pt.tensor.dscalar('dt')
tau_01 = pt.tensor.dscalar('tau_01')
tau_10 = pt.tensor.dscalar('tau_10')
a_01 = pt.tensor.dscalar('a_01')
a_10 = pt.tensor.dscalar('a_10')
# outputs_info
p = pt.tensor.dvector('p')

# sequences
times_value = np.arange(0, 100+1, 1)
# non_sequences
dt_value = 1
tau_01_value = 40
tau_10_value = 20
a_01_value = 5
a_10_value = 50
# outputs_info
p_value = np.array([1, 0]) # [state-0, state-1]

# Set up scan and evaluate
output, updates = pt.scan(fn=step, 
                          sequences=[times],
                          non_sequences=[dt, tau_01, tau_10, a_01, a_10],
                          outputs_info=[p])

f = pt.function(inputs=[t, 
                        dt, tau_01, tau_10, a_01, a_10,
                        p], 
                outputs=output,
                updates=updates)

result = f(times_value, 
           dt_value, tau_01_value, tau_10_value, a_01_value, a_10,
           p_value)

It seems this error is saying that my initial state p_value for the outputs_info p is 2-dimensional, though p_value is specified as 1-dimensional. The error also seems to be saying that p_new is 2-dimensional, though I was expecting that to be 1-dimensional.

What might I be missing here?

4

As a side note–though I am still interested in resolving my issues with a scan implementation–, I was able to vectorize the graph as you suggested.

Here is the calculation of a single step, taking in a single time t (same as above):

# sequences
times = pt.tensor.dvector('times')
t = pt.tensor.dscalar('t') # entry from `times` sequence
# non_sequences
dt = pt.tensor.dscalar('dt')
tau_01 = pt.tensor.dscalar('tau_01')
tau_10 = pt.tensor.dscalar('tau_10')
a_01 = pt.tensor.dscalar('a_01')
a_10 = pt.tensor.dscalar('a_10')
# outputs_info
p = pt.tensor.dvector('p')

# sequences
times_value = np.arange(0, 100+1, 1)
t_value = times_value[0] # first entry from `times` sequence
# non_sequences
dt_value = 1
tau_01_value = 40
tau_10_value = 20
a_01_value = 5
a_10_value = 50
# outputs_info
p_value = np.array([1, 0]) # [state-0, state-1]

# Computation of single step
gamma_01 = a_01 * pt.tensor.math.exp(-t/tau_01)
gamma_10 = a_10 * pt.tensor.math.exp(-t/tau_10)
T = pt.tensor.stacklists(
          [[-gamma_01,  gamma_10],
           [ gamma_01, -gamma_10]]
     ) * dt
p_new = pt.tensor.dot(T,p) + p

# Evaluate
p_new.eval({t:t_value, 
            dt:dt_value, tau_01:tau_01_value, tau_10:tau_10_value, a_01:a_01_value, a_10:a_10_value,
            p:p_value})


>> array([-4.,  5.])

And here is the vectorization:

p_new_vectorized = pt.graph.vectorize_graph([p_new], replace={t:times})
fn_vectorized = pt.function([times, 
                             dt, tau_01, tau_10, a_01, a_10, 
                             p], 
                             p_new_vectorized)
fn_vectorized(times_value, 
              dt_value, tau_01_value, tau_10_value, a_01_value, a_10_value, 
              p_value)


>> [array([[-4.        ,  5.        ],
           [-3.87654956,  4.87654956],
           [-3.75614712,  4.75614712],
           [-3.63871743,  4.63871743],
           ...
   ])]

EDIT: Upon further inspection, it seems this vectorization doesn’t accomplish the cumulative multiplication that I am seeking. That is, in this vectorization, the p_new output from step i doesn’t get used as the p input to step i+1.

1 Like
5

OK, I have gotten the scan implementation working. There were several issues in my previous attempt at a scan implementation. Those issue were:

  • a typo whereby result = f(..., a_10, ...) should have read result = f(..., a_10_value, ...)
  • a mis-ordering of arguments to the step function, which appeared in that function’s definition, and in the related calls to pt.function and the f function. Per scan tutorial e.g. 3, the recurrent outputs/inputs should appear “after the inputs that correspond to sequences but before those that correspond to non-sequences”.
  • an incorrect use of a scalar entry t to times in f = pt.function(inputs=[t, ... which should have been left as f = pt.function(inputs=[times, ... (though each call to step accesses a scalar entry t from the times sequence).

Here’s the latest attempt at a scan implementation of my desired routine, which runs without error:

def step(t, 
         p,
         dt, tau_01, tau_10, a_01, a_10):
     """
     step(sequences,
          outputs_info,
          non_sequences)
     "The step() function needs to expect one additional input 
     for each simple recurrent output. These inputs correspond to outputs 
     from previous iteration and are always after the inputs that 
     correspond to sequences but before those that correspond to non-sequences. 
     The[y] are received by the step() function in the order in which the 
     recurrent outputs are declared in the outputs_info sequence."
     """
     # Computation of single step
     gamma_01 = a_01 * pt.tensor.math.exp(-t/tau_01)
     gamma_10 = a_10 * pt.tensor.math.exp(-t/tau_10)
     T = pt.tensor.stacklists(
               [[-gamma_01,  gamma_10],
               [ gamma_01, -gamma_10]]
          ) * dt
     p_new = pt.tensor.dot(T,p) + p

     return p_new

# sequences
times = pt.tensor.dvector('times')
# non_sequences
dt = pt.tensor.dscalar('dt')
tau_01 = pt.tensor.dscalar('tau_01')
tau_10 = pt.tensor.dscalar('tau_10')
a_01 = pt.tensor.dscalar('a_01')
a_10 = pt.tensor.dscalar('a_10')
# outputs_info
p = pt.tensor.dvector('p')

# sequences
times_value = np.arange(0, 100+1, 1)
# non_sequences
dt_value = 1
tau_01_value = 40
tau_10_value = 20
a_01_value = 5
a_10_value = 50
# outputs_info
p_value = np.array([1, 0]) # [state-0, state-1]

# Set up scan and evaluate
output, updates = pt.scan(fn=step, 
                          sequences=[times],
                          non_sequences=[dt, tau_01, tau_10, a_01, a_10],
                          outputs_info=[p])

f = pt.function(inputs=[times, 
                        p,
                        dt, tau_01, tau_10, a_01, a_10], 
                outputs=output,
                updates=updates,
                on_unused_input='warn')

result = f(times_value, 
           p_value,
           dt_value, tau_01_value, tau_10_value, a_01_value, a_10_value)
2 Likes