Transform a matrix made of binomial vectors to ranges for consecutive zeros

Question

I am trying to figure out how to do this transformation symbolically in theano a matrix of undetermined size

From:

 [[0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 1],
  [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1],
  .
  .
  ]

To:

[[1, 2, 3, 0, 1, 2, 3, 4, 5, 0, 0, 1, 0, 1, 2, 3, 0],
 [1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 0, 0, 0, 0, 0, 0],
 .
 .
 ]

So for every consecutive 0 I want an increasing range and whenever I stumble on a 1 the range resets.

Answer 1

Here's one way to do it, using inefficient scans:

import theano
import theano.tensor as tt


def inner_step(x_t_t, y_t_tm1):
    return tt.switch(x_t_t, 0, y_t_tm1 + 1)


def outer_step(x_t):
    return theano.scan(inner_step, sequences=[x_t], outputs_info=[0])[0]


def compile():
    x = tt.bmatrix()
    y = theano.scan(outer_step, sequences=[x])[0]
    return theano.function([x], y)


def main():
    f = compile()
    data = [[0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 1],
            [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1]]
    print f(data)

main()

When run, this prints:

[[1 2 3 0 1 2 3 4 5 0 0 1 0 1 2 3 0]
 [1 2 3 4 5 6 7 8 0 1 2 0 0 0 0 0 0]]