Merge pull request #221 from mrocklin/roll_function

theano/tensor/basic.py

@@ -4290,7 +4290,33 @@ def join(axis, *tensors):
 pprint.assign(lambda pstate, r: r.owner and isinstance(r.owner.op, Join),
               printing.FunctionPrinter('join'))
 
+def roll(x, shift, axis=0):
+    """
+    Convenience function to roll `TensorType`s along the given axis.
+    Syntax copies numpy.roll function
+
+    Parameters
+    ----------
+     - x : a Tensor
+     - shift : int (symbolic or literal)
+        The number of places by which elements are shifted
+     - axis : int (symbolic or literal) (optional)
+        The axis along which elements are shifted.
+        Defaults to zero (deviation from numpy behavior)
+    """
 
+    # A slice of all elements in a dimension ':'
+    allslice = slice(None)
+    # List of slices describing the front half [:, :, shift:, :]
+    front_slice = slice(-shift, None)
+    front_list = ([allslice] * axis + [front_slice] +
+                  [allslice] * (x.ndim - axis - 1))
+    # List of slices describing the back half [:, :, :shift, :]
+    end_slice = slice(0, -shift)
+    end_list = [allslice]*axis + [end_slice] + [allslice]*(x.ndim-axis-1)
+    return join(axis,
+                Subtensor(front_list)(x),
+                Subtensor(end_list)(x))
 
 @constructor
 def shape_padleft(t, n_ones=1):

theano/tensor/tests/test_basic.py

@@ -30,7 +30,7 @@ from theano.tensor import (_shared, wvector, bvector, autocast_float_as,
         var, value, Join, shape, MaxAndArgmax, lscalar, zvector, exp,
         get_constant_value, ivector, reshape, scalar_from_tensor, scal,
         iscalars, arange,  dscalars, fvector, imatrix, numeric_grad,
-        opt, ComplexError, TensorDot, lvector, true_div, max, min, Split)
+        opt, ComplexError, TensorDot, lvector, true_div, max, min, Split, roll)
 from theano.tests import unittest_tools as utt
 
 
@@ -2813,6 +2813,35 @@ class T_Join_and_Split(unittest.TestCase):
         out = self.eval_outputs_and_check_join([s])
         self.assertTrue((out == want).all())
 
+    def test_roll(self):
+
+        # Test simple 1D example
+        a = self.shared(numpy.array([1, 2, 3, 4, 5, 6]))
+        b = roll(a, 2)
+        want = numpy.array([5, 6, 1, 2, 3, 4])
+        out = theano.function([], b)()
+
+        assert (out == want).all()
+
+        # Test 2D example - ensure that behavior matches numpy.roll behavior
+        a = self.shared(numpy.arange(21).reshape((3, 7)))
+        b = roll(a, -2, 1)
+
+        want = numpy.arange(21).reshape((3, 7))
+        want = numpy.roll(want, -2, 1)
+        out = theano.function([], b)()
+
+        assert (out == want).all()
+
+        # Test rolling on axis 0
+        want = numpy.arange(21).reshape((3, 7))
+        want = numpy.roll(want, -2, 0)
+        b = roll(a, -2, 0)
+        out = theano.function([], b)()
+
+        assert (out == want).all()
+
+
     def test_stack_vector(self):
         a = self.shared(numpy.array([1, 2, 3], dtype=self.floatX))
         b = as_tensor_variable(numpy.array([7, 8, 9], dtype=self.floatX))