Weighted matrix given weight-vector with scan (theano)

Question

I'm new to theano and still wrapping my head around scan. I want to compute a weighted matrix from row-weights, weight that by the weight's probability and get the corresponding weighted matrix. However, I'm having trouble iterating over weights in theano while keeping track of the matrix sum.

As a toy example:

import numpy as np

mtx = np.asarray([[1,0],[0,1],[0.5,0.5]]) 
weights = np.asarray([0.1,0.8]) #weight 1 and weight 2
weights_p = np.asarray([0.8, 0.2]) #prob. of weight 1 and weight 2

would be

weights_p[0] * (mtx * [weights[0],(1-weights[0])]) +\
weights_p[1] * (mtx * [weights[1],(1-weights[1])])

Exemplified more generally using numpy, indexing, and a for-loop my desired function would do this:

def get_weighted(mtx,weights,weights_p):
    mtx_store = np.zeros(np.shape(mtx))
    for idx in xrange(len(weights)):
        mtx_store += weights_p[idx] * (mtx * [weights[idx], 1-weights[idx]])
    return mtx_store

Now I need to do this in theano. What I tried:

import theano as t
v,w = t.tensor.vectors('v','w')
m,n = t.tensor.matrices('m','n')

def step(v, w, m, cum_sum):
    return v * (m * [w,1-w]) + cum_sum

output, updates = t.scan(fn=step,
                         sequences=[v,w],
                         non_sequences=[m],
                         outputs_info=[n])

get_weighted = t.function(inputs=[v,w,m,n],
                  outputs=output,
                  updates=updates)

My idea was to have an empty array to iteratively store the sum:

mtx_store = np.zeros(np.shape(mtx))
get_weighted(weights_p, weights, mtx, mtx_store)

But I'm getting:

array([[[ 1.  ,  0.  ],
        [ 0.  ,  1.  ],
        [ 0.5 ,  0.5 ]],

       [[ 1.16,  0.  ],
        [ 0.  ,  1.04],
       [ 0.58,  0.52]]])

Instead of

array([[ 0.24,  0.  ],
       [ 0.  ,  0.76],
       [ 0.12,  0.38]])

I'm sure this stems from my ill understanding of scan. What is wrong and how could it be done more efficiently?

Answer 1

I found the problem. For posterity: The main problem was that the syntax of scan wants:

sequences (if any), prior result(s) (if needed), non-sequences (if any)

whereas I had provided the arguments in this order:

sequences, non-sequences, prior-results

The correct code is as follows:

def step(v, w, cum_sum,m):
    return v * (m * [w,1-w]) + cum_sum

output, updates = t.scan(fn=step,
                         sequences=[v,w],
                         non_sequences=[m],
                         outputs_info=[t.tensor.zeros_like(m)])

final_result = output[-1] #take the final outcome of the sum


get_weighted = t.function(inputs=[v,w,m],
                  outputs=final_result,
                  updates=updates)

(Passing the matrix to store the arguments is apparently also not necessary. I don't think that this was the problem, but it can be directly specified as done in 'outputs_info' above)