roll using generalized elemwise op

test_roll.py

from random import randint
from time import time
import numpy as np
import theano as th
import theano.tensor as T
from theano.tensor.padding import idx, at_idx

N = 256

x = T.matrix()
shift = T.iscalar()

iy = idx(x, 0)
ix = idx(x, 1)
ny = T.shape(x)[0]
nx = T.shape(x)[1]

y_old = T.roll(x, shift)
y_new = at_idx(x, (iy+shift)%ny, ix)

fn_old = th.function([x, shift], y_old)
fn_new = th.function([x, shift], y_new)

xval = np.arange(N**2, dtype='float32').reshape(N, N)

# exlude the first call
fn_old(xval, 0)
fn_new(xval, 0)

beg = time()
for _ in range(100):
    sval = randint(0, N-1)
    yval1 = fn_old(xval, sval)
t_old = time() - beg

beg = time()
for _ in range(100):
    sval = randint(0, N-1)
    yval1 = fn_new(xval, sval)
t_new = time() - beg

print('old: %f new: %f' % (t_old, t_new))