PEP8 on new tensordot and its test.

theano/sandbox/cuda/basic_ops.py

@@ -2059,32 +2059,44 @@ class GpuContiguous(Op):
 
 gpu_contiguous = GpuContiguous()
 
+
 def tensordot(a, b, axes=2):
     """
-    implementation of tensordot that reduces to a regular matrix product. This allows tensordot to be GPU accelerated,
-    which isn't possible with the default Theano implementation (which is just a wrapper around numpy.tensordot).
-    based on code from Tijmen Tieleman's gnumpy http://www.cs.toronto.edu/~tijmen/gnumpy.html
+    Implementation of tensordot that reduces to a regular matrix product.
+
+    This allows tensordot to be GPU accelerated, which isn't possible
+    with the default Theano implementation (which is just a wrapper
+    around numpy.tensordot). based on code from Tijmen Tieleman's gnumpy
+    http://www.cs.toronto.edu/~tijmen/gnumpy.html
     """
     if numpy.isscalar(axes):
         # if 'axes' is a number of axes to multiply and sum over (trailing axes
         # of a, leading axes of b), we can just reshape and use dot.
-        outshape = tensor.concatenate([a.shape[:a.ndim - axes], b.shape[axes:]])
-        outndim = a.ndim + b.ndim - 2*axes
-        a_reshaped = a.reshape((tensor.prod(a.shape[:a.ndim - axes]), tensor.prod(a.shape[a.ndim - axes:])))
-        b_reshaped = b.reshape((tensor.prod(b.shape[:axes]), tensor.prod(b.shape[axes:])))
-        return tensor.dot(a_reshaped, b_reshaped).reshape(outshape, ndim=outndim)
+        outshape = tensor.concatenate([a.shape[:a.ndim - axes],
+                                      b.shape[axes:]])
+        outndim = a.ndim + b.ndim - (2 * axes)
+        a_reshaped = a.reshape((tensor.prod(a.shape[:a.ndim - axes]),
+                                tensor.prod(a.shape[a.ndim - axes:])))
+        b_reshaped = b.reshape((tensor.prod(b.shape[:axes]),
+                                tensor.prod(b.shape[axes:])))
+        return tensor.dot(a_reshaped, b_reshaped).reshape(outshape,
+                                                          ndim=outndim)
     elif len(axes) == 2:
         # if 'axes' is a pair of axis lists, we first shuffle the axes of a and
         # b to reduce this to the first case (note the recursion).
         a_other, b_other = tuple(axes[0]), tuple(axes[1])
         num_axes = len(a_other)
-        a_order = tuple(x for x in tuple(xrange(a.ndim)) if x not in a_other) + a_other
-        b_order = b_other + tuple(x for x in tuple(xrange(b.ndim)) if x not in b_other)
+        a_order = (tuple(x for x in tuple(xrange(a.ndim)) if x not in a_other)
+                + a_other)
+        b_order = (b_other
+                + tuple(x for x in tuple(xrange(b.ndim)) if x not in b_other))
         a_shuffled = a.dimshuffle(a_order)
         b_shuffled = b.dimshuffle(b_order)
         return tensordot(a_shuffled, b_shuffled, num_axes)
     else:
-        raise ValueError("Axes should be scalar valued or a list/tuple of len 2.")
+        raise ValueError(
+            "Axes should be scalar valued or a list/tuple of len 2.",
+            axes)
 
 # Those are predifined CudaNdarrayType as done in tensor.basic
 # Useful mostly for test as the gpu op are inserted automatically...

theano/sandbox/cuda/tests/test_basic_ops.py

@@ -870,16 +870,23 @@ def test_shared_cudandarray():
     a = cuda.shared_constructor(cuda.CudaNdarray.zeros((2,3)))
     assert isinstance(a.type, tcn.CudaNdarrayType)
 
+
 def test_tensordot_reshape():
-    '''Test that the tensordot implementation using dimshuffle, reshape and dot
-    gives the same results as the default (numpy) version'''
+    '''Test alternative tensordot implementation.
+
+    Test that the tensordot implementation using dimshuffle, reshape and dot
+    gives the same results as the default (numpy) version.
+    '''
     # define some tensors
     a = numpy.arange(20, dtype=theano.config.floatX) / 20.0
     b = numpy.arange(10, dtype=theano.config.floatX) / 10.0
     c = numpy.arange(5, dtype=theano.config.floatX) / 5.0
     d = numpy.arange(8, dtype=theano.config.floatX) / 8.0
 
-    tensor1 = numpy.tensordot(a, numpy.tensordot(b, numpy.tensordot(c, d, 0), 0), 0)
+    tensor1 = numpy.tensordot(
+            a,
+            numpy.tensordot(b, numpy.tensordot(c, d, 0), 0),
+            0)
     tensor2 = numpy.tensordot(c, numpy.tensordot(d, a, 0), 0)
     tensor3 = tensor2.swapaxes(1, 2).swapaxes(0, 2)  # d, a, c
 
@@ -887,17 +894,29 @@ def test_tensordot_reshape():
     y = T.tensor3('y')
 
     # case 1: number of axes to sum over
-    default1 = theano.function([x,y], T.tensordot(x, y, 2))(tensor1, tensor2)
-    reshape1 = theano.function([x,y], B.tensordot(x, y, 2))(tensor1, tensor2)
+    default1 = theano.function([x, y], T.tensordot(x, y, 2))(tensor1, tensor2)
+    reshape1 = theano.function([x, y], B.tensordot(x, y, 2))(tensor1, tensor2)
     assert numpy.allclose(default1, reshape1)
 
     # case 2: axis pairs
-    default2 = theano.function([x,y], T.tensordot(x, y, axes=[(0, 3), (1, 0)]))(tensor1, tensor3)
-    reshape2 = theano.function([x,y], B.tensordot(x, y, axes=[(0, 3), (1, 0)]))(tensor1, tensor3)
+    default2 = theano.function(
+            [x, y],
+            T.tensordot(x, y, axes=[(0, 3), (1, 0)])
+            )(tensor1, tensor3)
+    reshape2 = theano.function(
+            [x, y],
+            B.tensordot(x, y, axes=[(0, 3), (1, 0)])
+            )(tensor1, tensor3)
     assert numpy.allclose(default2, reshape2)
 
-    default3 = theano.function([x,y], T.tensordot(x, y, axes=[(0, 3, 2), (1, 0, 2)]))(tensor1, tensor3)
-    reshape3 = theano.function([x,y], B.tensordot(x, y, axes=[(0, 3, 2), (1, 0, 2)]))(tensor1, tensor3)
+    default3 = theano.function(
+            [x, y],
+            T.tensordot(x, y, axes=[(0, 3, 2), (1, 0, 2)])
+            )(tensor1, tensor3)
+    reshape3 = theano.function(
+            [x, y],
+            B.tensordot(x, y, axes=[(0, 3, 2), (1, 0, 2)])
+            )(tensor1, tensor3)
     assert numpy.allclose(default3, reshape3)