Repository: lisa-lab/pylearn2

Issue:
C01B sampling tests for pylearn2.expr.probabilistic_max_pooling
There are already tests for functionality other than sampling and sampling in formats other than C01B. Just copy-paste some of these to test the C01B sampling.



Discussion:

Resolving patch:
diff --git a/pylearn2/expr/probabilistic_max_pooling.py b/pylearn2/expr/probabilistic_max_pooling.py
--- a/pylearn2/expr/probabilistic_max_pooling.py
+++ b/pylearn2/expr/probabilistic_max_pooling.py
@@ -33,7 +33,8 @@
 from theano.gof.op import get_debug_values
 import warnings
 
-def max_pool(z, pool_shape, top_down = None, theano_rng = None):
+
+def max_pool(z, pool_shape, top_down=None, theano_rng=None):
     """
     .. todo::
 
@@ -54,71 +55,78 @@ def max_pool(z, pool_shape, top_down = None, theano_rng = None):
 
     all 4-tensors are formatted with axes ('b', 'c', 0, 1).
     This is for maximum speed when using theano's conv2d
-    to generate z and top_down, or when using it to infer conditionals of other layers
-    using the return values.
+    to generate z and top_down, or when using it to infer conditionals of
+    other layers using the return values.
 
     Detailed description:
 
-    Suppose you have a variable h that lives in a Conv2DSpace h_space and you want to
-    pool it down to a variable p that lives in a smaller Conv2DSpace p.
+    Suppose you have a variable h that lives in a Conv2DSpace h_space and
+    you want to pool it down to a variable p that lives in a smaller
+    Conv2DSpace p.
 
     This function does that, using non-overlapping pools.
 
-    Specifically, consider one channel of h. h must have a height that is a multiple of
-    pool_shape[0] and a width that is a multiple of pool_shape[1]. A channel of h can
-    thus be broken down into non-overlapping rectangles of shape pool_shape.
+    Specifically, consider one channel of h. h must have a height that is a
+    multiple of pool_shape[0] and a width that is a multiple of pool_shape[1].
+    A channel of h can thus be broken down into non-overlapping rectangles
+    of shape pool_shape.
 
     Now consider one rectangular pooled region within one channel of h.
-    I now use 'h' to refer just to this rectangle, and 'p' to refer to just the one pooling
-    unit associated with that rectangle.
-    We assume that the space that h and p live in is constrained such that h and p are
-    both binary and p = max(h). To reduce the state-space in order to make probabilistic
-    computations cheaper we also constrain sum(h) <= 1.
-    Suppose h contains k different units. Suppose that the only term in the model's energy
-    function involving h is -(z*h).sum() (elemwise multiplication) and the only term in
+    I now use 'h' to refer just to this rectangle, and 'p' to refer to
+    just the one pooling unit associated with that rectangle.
+    We assume that the space that h and p live in is constrained such
+    that h and p are both binary and p = max(h). To reduce the state-space
+    in order to make probabilistic computations cheaper we also
+    constrain sum(h) <= 1.
+    Suppose h contains k different units. Suppose that the only term
+    in the model's energy function involving h is -(z*h).sum()
+    (elemwise multiplication) and the only term in
     the model's energy function involving p is -(top_down*p).sum().
 
     Then P(h[i] = 1) = softmax( [ z[1], z[2], ..., z[k], -top_down] )[i]
     and P(p = 1) = 1-softmax( [z[1], z[2], ..., z[k], -top_down])[k]
 
 
-    This variation of the function assumes that z, top_down, and all return values use
-    Conv2D axes ('b', 'c', 0, 1).
-    This variation of the function implements the softmax using a theano graph of exp,
-    maximum, sub, and div operations.
+    This variation of the function assumes that z, top_down, and all
+    return values use Conv2D axes ('b', 'c', 0, 1).
+    This variation of the function implements the softmax using a
+    theano graph of exp, maximum, sub, and div operations.
 
     Performance notes:
-    It might be possible to make a faster implementation with different theano ops.
-    rather than using set_subtensor, it might be possible to use the stuff in
-    theano.sandbox.neighbours. Probably not possible, or at least nasty, because
-    that code isn't written with multiple channels in mind, and I don't think
-    just a reshape can fix it. Some work on this in galatea.cond.neighbs.py
-    At some point images2neighbs' gradient was broken so check that it has been fixed
-    before sinking too much time into this.
-
-    Stabilizing the softmax is also another source of slowness. Here it is stabilized
-    with several calls to maximum and sub. It might also be possible to stabilize it
-    with T.maximum(-top_down,T.signal.downsample.max_pool(z)). Don't know if that would
-    be faster or slower.
-
-    Elsewhere in this file I implemented the softmax with a reshape and call to Softmax /
-    SoftmaxWithBias. This is slower, even though Softmax is faster on the GPU than the
+    It might be possible to make a faster implementation with different
+    theano ops. rather than using set_subtensor, it might be possible
+    to use the stuff in theano.sandbox.neighbours. Probably not possible,
+    or at least nasty, because that code isn't written with multiple
+    channels in mind, and I don't think just a reshape can fix it.
+    Some work on this in galatea.cond.neighbs.py
+    At some point images2neighbs' gradient was broken so check that
+    it has been fixed before sinking too much time into this.
+
+    Stabilizing the softmax is also another source of slowness.
+    Here it is stabilized with several calls to maximum and sub.
+    It might also be possible to stabilize it with
+    T.maximum(-top_down,T.signal.downsample.max_pool(z)).
+    Don't know if that would be faster or slower.
+
+    Elsewhere in this file I implemented the softmax with a reshape
+    and call to Softmax / SoftmaxWithBias.
+    This is slower, even though Softmax is faster on the GPU than the
     equivalent max/sub/exp/div graph. Maybe the reshape is too expensive.
 
-    Benchmarks show that most of the time is spent in GpuIncSubtensor when running on
-    gpu. So it is mostly that which needs a faster implementation.
-    One other way to implement this would be with a linear.Conv2D.lmul_T, where the
-    convolution stride is equal to the pool width, and the thing to multiply with is
-    the hparts stacked along the channel axis. Unfortunately, conv2D doesn't work right
-    with stride > 2 and is pretty slow for stride 2. Conv3D is used to mitigat some
-    of this, but only has CPU code.
+    Benchmarks show that most of the time is spent in GpuIncSubtensor
+    when running on gpu. So it is mostly that which needs a faster
+    implementation. One other way to implement this would be with
+    a linear.Conv2D.lmul_T, where the convolution stride is equal to
+    the pool width, and the thing to multiply with is the hparts stacked
+    along the channel axis. Unfortunately, conv2D doesn't work right
+    with stride > 2 and is pretty slow for stride 2. Conv3D is used to
+    mitigat some of this, but only has CPU code.
     """
 
     z_name = z.name
     if z_name is None:
         z_name = 'anon_z'
 
-
     batch_size, ch, zr, zc = z.shape
 
     r, c = pool_shape
@@ -136,21 +144,21 @@ def max_pool(z, pool_shape, top_down = None, theano_rng = None):
     for i in xrange(r):
         zpart.append([])
         for j in xrange(c):
-            cur_part = z[:,:,i:zr:r,j:zc:c]
+            cur_part = z[:, :, i:zr:r, j:zc:c]
             if z_name is not None:
-                cur_part.name = z_name + '[%d,%d]' % (i,j)
-            zpart[i].append( cur_part )
+                cur_part.name = z_name + '[%d,%d]' % (i, j)
+            zpart[i].append(cur_part)
             if mx is None:
                 mx = T.maximum(t, cur_part)
                 if cur_part.name is not None:
-                    mx.name = 'max(-top_down,'+cur_part.name+')'
+                    mx.name = 'max(-top_down,' + cur_part.name + ')'
             else:
                 max_name = None
                 if cur_part.name is not None:
-                    mx_name = 'max('+cur_part.name+','+mx.name+')'
-                mx = T.maximum(mx,cur_part)
+                    mx_name = 'max(' + cur_part.name + ',' + mx.name + ')'
+                mx = T.maximum(mx, cur_part)
                 mx.name = mx_name
-    mx.name = 'local_max('+z_name+')'
+    mx.name = 'local_max(' + z_name + ')'
 
     pt = []
 
@@ -162,7 +170,7 @@ def max_pool(z, pool_shape, top_down = None, theano_rng = None):
             safe.name = 'safe_z(%s)' % z_ij.name
             cur_pt = T.exp(safe)
             cur_pt.name = 'pt(%s)' % z_ij.name
-            pt[-1].append( cur_pt )
+            pt[-1].append(cur_pt)
 
     off_pt = T.exp(t - mx)
     off_pt.name = 'p_tilde_off(%s)' % z_name
@@ -179,14 +187,14 @@ def max_pool(z, pool_shape, top_down = None, theano_rng = None):
 
     hpart = []
     for i in xrange(r):
-        hpart.append( [ pt_ij / denom for pt_ij in pt[i] ] )
+        hpart.append([pt_ij / denom for pt_ij in pt[i]])
 
     h = T.alloc(0., batch_size, ch, zr, zc)
 
     for i in xrange(r):
         for j in xrange(c):
             h.name = 'h_interm'
-            h = T.set_subtensor(h[:,:,i:zr:r,j:zc:c],hpart[i][j])
+            h = T.set_subtensor(h[:, :, i:zr:r, j:zc:c], hpart[i][j])
 
     h.name = 'h(%s)' % z_name
 
@@ -199,42 +207,48 @@ def max_pool(z, pool_shape, top_down = None, theano_rng = None):
                 events.append(hpart[i][j])
         events.append(off_prob)
 
-        events = [ event.dimshuffle(0,1,2,3,'x') for event in events ]
+        events = [event.dimshuffle(0, 1, 2, 3, 'x') for event in events]
 
         events = tuple(events)
 
-        stacked_events = T.concatenate( events, axis = 4)
+        stacked_events = T.concatenate(events, axis=4)
 
         rows = zr // pool_shape[0]
         cols = zc // pool_shape[1]
         outcomes = pool_shape[0] * pool_shape[1] + 1
         assert stacked_events.ndim == 5
-        for se, bs, r, c, chv in get_debug_values(stacked_events, batch_size, rows, cols, ch):
+        for se, bs, r, c, chv in get_debug_values(stacked_events, batch_size,
+                                                  rows, cols, ch):
             assert se.shape[0] == bs
             assert se.shape[1] == r
             assert se.shape[2] == c
             assert se.shape[3] == chv
             assert se.shape[4] == outcomes
-        reshaped_events = stacked_events.reshape((batch_size * rows * cols * ch, outcomes))
+        reshaped_events = stacked_events.reshape((
+            batch_size * rows * cols * ch, outcomes))
 
-        multinomial = theano_rng.multinomial(pvals = reshaped_events, dtype = p.dtype)
+        multinomial = theano_rng.multinomial(pvals=reshaped_events,
+                                             dtype=p.dtype)
 
-        reshaped_multinomial = multinomial.reshape((batch_size, ch, rows, cols, outcomes))
+        reshaped_multinomial = multinomial.reshape((batch_size, ch, rows,
+                                                    cols, outcomes))
 
         h_sample = T.alloc(0., batch_size, ch, zr, zc)
 
         idx = 0
         for i in xrange(r):
             for j in xrange(c):
-                h_sample = T.set_subtensor(h_sample[:,:,i:zr:r,j:zc:c],
-                        reshaped_multinomial[:,:,:,:,idx])
+                h_sample = T.set_subtensor(h_sample[:, :, i:zr:r, j:zc:c],
+                                           reshaped_multinomial[:, :, :, :,
+                                           idx])
                 idx += 1
 
-        p_sample = 1 - reshaped_multinomial[:,:,:,:,-1]
+        p_sample = 1 - reshaped_multinomial[:, :, :, :, -1]
 
         return p, h, p_sample, h_sample
 
-def max_pool_c01b(z, pool_shape, top_down = None, theano_rng = None):
+
+def max_pool_c01b(z, pool_shape, top_down=None, theano_rng=None):
     """
     .. todo::
 
@@ -245,16 +259,17 @@ def max_pool_c01b(z, pool_shape, top_down = None, theano_rng = None):
 
     Performance notes:
     Stabilizing the softmax is one source slowness. Here it is stabilized
-    with several calls to maximum and sub. It might also be possible to stabilize it
-    with T.maximum(-top_down,<cuda convnet max pooling>). Don't know if that would
-    be faster or slower.
-
-    Benchmarks show that most of the time is spent in GpuIncSubtensor when running on
-    gpu. So it is mostly that which needs a faster implementation.
-    One other way to implement this would be with cuda convnet convolution, where the
-    convolution stride is equal to the pool width, and the thing to multiply with is
-    the hparts stacked along the channel axis. This isn't a feasible solution for
-    max_pool because of theano convolution's poor support for strides, but for cuda
+    with several calls to maximum and sub. It might also be possible to
+    stabilize it with T.maximum(-top_down,<cuda convnet max pooling>).
+    Don't know if that would be faster or slower.
+
+    Benchmarks show that most of the time is spent in GpuIncSubtensor
+    when running on gpu. So it is mostly that which needs a faster
+    implementation. One other way to implement this would be with cuda
+    convnet convolution, where the convolution stride is equal to the
+    pool width, and the thing to multiply with is the hparts stacked
+    along the channel axis. This isn't a feasible solution for max_pool
+    because of theano convolution's poor support for strides, but for cuda
     convnet it could give a speedup.
     """
 
@@ -262,7 +277,6 @@ def max_pool_c01b(z, pool_shape, top_down = None, theano_rng = None):
     if z_name is None:
         z_name = 'anon_z'
 
-
     ch, zr, zc, batch_size = z.shape
 
     r, c = pool_shape
@@ -280,21 +294,21 @@ def max_pool_c01b(z, pool_shape, top_down = None, theano_rng = None):
     for i in xrange(r):
         zpart.append([])
         for j in xrange(c):
-            cur_part = z[:,i:zr:r,j:zc:c,:]
+            cur_part = z[:, i:zr:r, j:zc:c, :]
             if z_name is not None:
-                cur_part.name = z_name + '[%d,%d]' % (i,j)
-            zpart[i].append( cur_part )
+                cur_part.name = z_name + '[%d, %d]' % (i, j)
+            zpart[i].append(cur_part)
             if mx is None:
                 mx = T.maximum(t, cur_part)
                 if cur_part.name is not None:
-                    mx.name = 'max(-top_down,'+cur_part.name+')'
+                    mx.name = 'max(-top_down,' + cur_part.name + ')'
             else:
                 max_name = None
                 if cur_part.name is not None:
-                    mx_name = 'max('+cur_part.name+','+mx.name+')'
-                mx = T.maximum(mx,cur_part)
+                    mx_name = 'max(' + cur_part.name + ',' + mx.name + ')'
+                mx = T.maximum(mx, cur_part)
                 mx.name = mx_name
-    mx.name = 'local_max('+z_name+')'
+    mx.name = 'local_max(' + z_name + ')'
 
     pt = []
 
@@ -306,7 +320,7 @@ def max_pool_c01b(z, pool_shape, top_down = None, theano_rng = None):
             safe.name = 'safe_z(%s)' % z_ij.name
             cur_pt = T.exp(safe)
             cur_pt.name = 'pt(%s)' % z_ij.name
-            pt[-1].append( cur_pt )
+            pt[-1].append(cur_pt)
 
     off_pt = T.exp(t - mx)
     off_pt.name = 'p_tilde_off(%s)' % z_name
@@ -323,14 +337,14 @@ def max_pool_c01b(z, pool_shape, top_down = None, theano_rng = None):
 
     hpart = []
     for i in xrange(r):
-        hpart.append( [ pt_ij / denom for pt_ij in pt[i] ] )
+        hpart.append([pt_ij / denom for pt_ij in pt[i]])
 
     h = T.alloc(0., ch, zr, zc, batch_size)
 
     for i in xrange(r):
         for j in xrange(c):
             h.name = 'h_interm'
-            h = T.set_subtensor(h[:,i:zr:r,j:zc:c,:],hpart[i][j])
+            h = T.set_subtensor(h[:, i:zr:r, j:zc:c, :], hpart[i][j])
 
     h.name = 'h(%s)' % z_name
 
@@ -343,42 +357,38 @@ def max_pool_c01b(z, pool_shape, top_down = None, theano_rng = None):
                 events.append(hpart[i][j])
         events.append(off_prob)
 
-        events = [ event.dimshuffle(0,1,2,3,'x') for event in events ]
+        events = [event.dimshuffle(0, 1, 2, 3, 'x') for event in events]
 
         events = tuple(events)
 
-        stacked_events = T.concatenate( events, axis = 4)
+        stacked_events = T.concatenate(events, axis=4)
 
-        rows = zr // pool_shape[0]
-        cols = zc // pool_shape[1]
-        outcomes = pool_shape[0] * pool_shape[1] + 1
-        assert stacked_events.ndim == 5
-        for se, bs, r, c, chv in get_debug_values(stacked_events, batch_size, rows, cols, ch):
-            assert se.shape[0] == chv
-            assert se.shape[1] == r
-            assert se.shape[2] == c
-            assert se.shape[3] == bs
-            assert se.shape[4] == outcomes
-        reshaped_events = stacked_events.reshape((rows * cols * ch * batch_size, outcomes))
+        ch, rows, cols, batch_size, outcomes = stacked_events.shape
+        reshaped_events = stacked_events.reshape((ch * rows * cols *
+                                                  batch_size, outcomes))
 
-        multinomial = theano_rng.multinomial(pvals = reshaped_events, dtype = p.dtype)
+        multinomial = theano_rng.multinomial(pvals=reshaped_events,
+                                             dtype=p.dtype)
 
-        reshaped_multinomial = multinomial.reshape((ch, rows, cols, batch_size, outcomes))
+        reshaped_multinomial = multinomial.reshape((ch, rows, cols, batch_size,
+                                                    outcomes))
 
-        h_sample = T.alloc(0., ch, zr, zc, rows, cols, batch_size, outcomes)
+        h_sample = T.alloc(0., ch, zr, zc, batch_size)
 
         idx = 0
         for i in xrange(r):
             for j in xrange(c):
-                h_sample = T.set_subtensor(h_sample[:,i:zr:r,j:zc:c,:],
-                        reshaped_multinomial[:,:,:,:,idx])
+                h_sample = T.set_subtensor(h_sample[:, i:zr:r, j:zc:c, :],
+                                           reshaped_multinomial[:, :, :, :,
+                                           idx])
                 idx += 1
 
-        p_sample = 1 - reshaped_multinomial[:,:,:,:,-1]
+        p_sample = 1 - reshaped_multinomial[:, :, :, :, -1]
 
         return p, h, p_sample, h_sample
 
-def max_pool_channels(z, pool_size, top_down = None, theano_rng = None):
+
+def max_pool_channels(z, pool_size, top_down=None, theano_rng=None):
     """
     .. todo::
 
@@ -419,10 +429,13 @@ def max_pool_channels(z, pool_size, top_down = None, theano_rng = None):
             return p, h
         else:
             t1 = time.time()
-            p_samples = theano_rng.binomial(p=p, size=p.shape, dtype=p.dtype, n=1)
+            p_samples = theano_rng.binomial(p=p, size=p.shape,
+                                            dtype=p.dtype, n=1)
             t2 = time.time()
             if t2 - t1 > 0.5:
-                warnings.warn("TODO: speed up theano's random number seeding. max pooling spent "+str(t2-t1)+" in a call to theano_rng.binomial.")
+                warnings.warn("TODO: speed up theano's random number seeding. "
+                              "max pooling spent " + str(t2-t1) +
+                              "in a call to theano_rng.binomial.")
             h_samples = p_samples
             return p_samples, h_samples, p_samples, h_samples
     else:
@@ -438,21 +451,21 @@ def max_pool_channels(z, pool_size, top_down = None, theano_rng = None):
 
         zpart = []
         for i in xrange(pool_size):
-            cur_part = z[:,i:n:pool_size]
+            cur_part = z[:, i:n:pool_size]
             if z_name is not None:
                 cur_part.name = z_name + '[%d]' % (i)
-            zpart.append( cur_part )
+            zpart.append(cur_part)
             if mx is None:
                 mx = T.maximum(t, cur_part)
                 if cur_part.name is not None:
-                    mx.name = 'max(-top_down,'+cur_part.name+')'
+                    mx.name = 'max(-top_down,' + cur_part.name + ')'
             else:
                 max_name = None
                 if cur_part.name is not None:
-                    mx_name = 'max('+cur_part.name+','+mx.name+')'
-                mx = T.maximum(mx,cur_part)
+                    mx_name = 'max(' + cur_part.name + ',' + mx.name + ')'
+                mx = T.maximum(mx, cur_part)
                 mx.name = mx_name
-        mx.name = 'local_max('+z_name+')'
+        mx.name = 'local_max(' + z_name + ')'
 
         pt = []
 
@@ -463,7 +476,7 @@ def max_pool_channels(z, pool_size, top_down = None, theano_rng = None):
             cur_pt = T.exp(safe)
             cur_pt.name = 'pt(%s)' % z_i.name
             assert cur_pt.ndim == 2
-            pt.append( cur_pt )
+            pt.append(cur_pt)
 
         off_pt = T.exp(t - mx)
         assert off_pt.ndim == 2
@@ -479,21 +492,21 @@ def max_pool_channels(z, pool_size, top_down = None, theano_rng = None):
         p = 1. - off_prob
         assert p.dtype == z.dtype
 
-        hpart = [ pt_i / denom for pt_i in pt ]
+        hpart = [pt_i / denom for pt_i in pt]
 
         h = T.alloc(0., batch_size, n)
 
         for i in xrange(pool_size):
             h.name = 'h_interm'
             hp = hpart[i]
-            sub_h = h[:,i:n:pool_size]
+            sub_h = h[:, i:n:pool_size]
             assert sub_h.ndim == 2
             assert hp.ndim == 2
             for hv, hsv, hpartv in get_debug_values(h, sub_h, hp):
                 print hv.shape
                 print hsv.shape
                 print hpartv.shape
-            h = T.set_subtensor(sub_h,hp)
+            h = T.set_subtensor(sub_h, hp)
 
     p.name = 'p(%s)' % z_name
     h.name = 'h(%s)' % z_name
@@ -506,38 +519,45 @@ def max_pool_channels(z, pool_size, top_down = None, theano_rng = None):
             events.append(hpart[i])
         events.append(off_prob)
 
-        events = [ event.dimshuffle(0,1,'x') for event in events ]
+        events = [event.dimshuffle(0, 1, 'x') for event in events]
 
         events = tuple(events)
 
-        stacked_events = T.concatenate( events, axis = 2)
+        stacked_events = T.concatenate(events, axis=2)
 
         outcomes = pool_size + 1
-        reshaped_events = stacked_events.reshape((batch_size * n // pool_size, outcomes))
+        reshaped_events = stacked_events.reshape((batch_size * n // pool_size,
+                                                  outcomes))
 
         t1 = time.time()
-        multinomial = theano_rng.multinomial(pvals = reshaped_events, dtype = p.dtype)
+        multinomial = theano_rng.multinomial(pvals=reshaped_events,
+                                             dtype=p.dtype)
         t2 = time.time()
         if t2 - t1 > 0.5:
-            warnings.warn("TODO: speed up theano's random number seeding. max pooling spent "+str(t2-t1)+" in a call to theano_rng.multinomial.")
+            warnings.warn("TODO: speed up theano's random number seeding."
+                          "max pooling spent " + str(t2-t1) +
+                          " in a call to theano_rng.multinomial.")
 
-        reshaped_multinomial = multinomial.reshape((batch_size, n // pool_size, outcomes))
+        reshaped_multinomial = multinomial.reshape((batch_size,
+                                                    n // pool_size,
+                                                    outcomes))
 
         h_sample = T.zeros_like(z)
 
         idx = 0
         for i in xrange(pool_size):
-            h_sample = T.set_subtensor(h_sample[:,i:n:pool_size],
-                        reshaped_multinomial[:,:,idx])
+            h_sample = T.set_subtensor(h_sample[:, i:n:pool_size],
+                                       reshaped_multinomial[:, :, idx])
             idx += 1
 
-        p_sample = 1 - reshaped_multinomial[:,:,-1]
+        p_sample = 1 - reshaped_multinomial[:, :, -1]
 
         assert h_sample.dtype == z.dtype
 
         return p, h, p_sample, h_sample
 
-def max_pool_python(z, pool_shape, top_down = None):
+
+def max_pool_python(z, pool_shape, top_down=None):
     """
     .. todo::
 
@@ -555,25 +575,26 @@ def max_pool_python(z, pool_shape, top_down = None):
     assert zr % r == 0
     assert zc % c == 0
 
-    h = np.zeros(z.shape, dtype = z.dtype)
-    p = np.zeros( (batch_size, zr /r, zc /c, ch), dtype = z.dtype)
+    h = np.zeros(z.shape, dtype=z.dtype)
+    p = np.zeros((batch_size, zr / r, zc / c, ch), dtype=z.dtype)
     if top_down is None:
         top_down = p.copy()
 
-    for u in xrange(0,zr,r):
-        for l in xrange(0,zc,c):
-            pt = np.exp(z[:,u:u+r,l:l+c,:])
-            off_pt = np.exp(-top_down[:,u/r,l/c,:])
+    for u in xrange(0, zr, r):
+        for l in xrange(0, zc, c):
+            pt = np.exp(z[:, u:u+r, l:l+c, :])
+            off_pt = np.exp(-top_down[:, u/r, l/c, :])
             denom = pt.sum(axis=1).sum(axis=1) + off_pt
-            p[:,u/r,l/c,:] = 1. - off_pt / denom
+            p[:, u/r, l/c, :] = 1. - off_pt / denom
             for i in xrange(batch_size):
                 for j in xrange(ch):
-                    pt[i,:,:,j] /= denom[i,j]
-            h[:,u:u+r,l:l+c,:] = pt
+                    pt[i, :, :, j] /= denom[i, j]
+            h[:, u:u+r, l:l+c, :] = pt
 
     return p, h
 
-def max_pool_channels_python(z, pool_size, top_down = None):
+
+def max_pool_channels_python(z, pool_size, top_down=None):
     """
     .. todo::
 
@@ -588,23 +609,24 @@ def max_pool_channels_python(z, pool_size, top_down = None):
 
     assert n % pool_size == 0
 
-    h = np.zeros(z.shape, dtype = z.dtype)
-    p = np.zeros( (batch_size, n / pool_size), dtype = z.dtype)
+    h = np.zeros(z.shape, dtype=z.dtype)
+    p = np.zeros((batch_size, n / pool_size), dtype=z.dtype)
     if top_down is None:
         top_down = p.copy()
 
-    for i in xrange(0,n / pool_size):
-        pt = np.exp(z[:,i*pool_size:(i+1)*pool_size])
-        off_pt = np.exp(-top_down[:,i])
+    for i in xrange(0, n / pool_size):
+        pt = np.exp(z[:, i*pool_size:(i+1)*pool_size])
+        off_pt = np.exp(-top_down[:, i])
         denom = pt.sum(axis=1) + off_pt
         assert denom.ndim == 1
-        p[:,i] = 1. - off_pt / denom
+        p[:, i] = 1. - off_pt / denom
         for j in xrange(batch_size):
             for k in xrange(pool_size):
-                h[j,i*pool_size+k] = pt[j,k] / denom[j]
+                h[j, i*pool_size+k] = pt[j, k] / denom[j]
 
     return p, h
 
+
 def max_pool_unstable(z, pool_shape):
     """
     .. todo::
@@ -612,8 +634,8 @@ def max_pool_unstable(z, pool_shape):
         WRITEME properly
 
     A version of max_pool that does not numerically stabilize the softmax.
-    This is faster, but prone to both overflow and underflow in the intermediate
-    computations.
+    This is faster, but prone to both overflow and underflow in the
+    intermediate computations.
     Mostly useful for benchmarking, to determine how much speedup we could
     hope to get by using a better stabilization method.
     Also, this uses the ('b', 0, 1, 'c') format.
@@ -628,12 +650,12 @@ def max_pool_unstable(z, pool_shape):
     for i in xrange(r):
         zpart.append([])
         for j in xrange(c):
-            zpart[i].append( z[:,i:zr:r,j:zc:c,:] )
+            zpart[i].append(z[:, i:zr:r, j:zc:c, :])
 
     pt = []
 
     for i in xrange(r):
-        pt.append( [ T.exp(z_ij) for z_ij in zpart[i] ] )
+        pt.append([T.exp(z_ij) for z_ij in zpart[i]])
 
     denom = 1.
 
@@ -645,17 +667,18 @@ def max_pool_unstable(z, pool_shape):
 
     hpart = []
     for i in xrange(r):
-        hpart.append( [ pt_ij / denom for pt_ij in pt[i] ] )
+        hpart.append([pt_ij / denom for pt_ij in pt[i]])
 
     h = T.alloc(0., batch_size, zr, zc, ch)
 
     for i in xrange(r):
         for j in xrange(c):
-            h = T.set_subtensor(h[:,i:zr:r,j:zc:c,:],hpart[i][j])
+            h = T.set_subtensor(h[:, i:zr:r, j:zc:c, :], hpart[i][j])
 
     return p, h
 
-def max_pool_b01c(z, pool_shape, top_down = None, theano_rng = None):
+
+def max_pool_b01c(z, pool_shape, top_down=None, theano_rng=None):
     """
     .. todo::
 
@@ -685,19 +708,19 @@ def max_pool_b01c(z, pool_shape, top_down = None, theano_rng = None):
     for i in xrange(r):
         zpart.append([])
         for j in xrange(c):
-            cur_part = z[:,i:zr:r,j:zc:c,:]
+            cur_part = z[:, i:zr:r, j:zc:c, :]
             if z_name is not None:
-                cur_part.name = z_name + '[%d,%d]' % (i,j)
-            zpart[i].append( cur_part )
+                cur_part.name = z_name + '[%d, %d]' % (i, j)
+            zpart[i].append(cur_part)
             if mx is None:
                 mx = T.maximum(t, cur_part)
                 if cur_part.name is not None:
-                    mx.name = 'max(-top_down,'+cur_part.name+')'
+                    mx.name = 'max(-top_down,' + cur_part.name + ')'
             else:
                 max_name = None
                 if cur_part.name is not None:
-                    mx_name = 'max('+cur_part.name+','+mx.name+')'
-                mx = T.maximum(mx,cur_part)
+                    mx_name = 'max(' + cur_part.name + ','+mx.name + ')'
+                mx = T.maximum(mx, cur_part)
                 mx.name = mx_name
     mx.name = 'local_max('+z_name+')'
 
@@ -711,7 +734,7 @@ def max_pool_b01c(z, pool_shape, top_down = None, theano_rng = None):
             safe.name = 'safe_z(%s)' % z_ij.name
             cur_pt = T.exp(safe)
             cur_pt.name = 'pt(%s)' % z_ij.name
-            pt[-1].append( cur_pt )
+            pt[-1].append(cur_pt)
 
     off_pt = T.exp(t - mx)
     off_pt.name = 'p_tilde_off(%s)' % z_name
@@ -728,13 +751,13 @@ def max_pool_b01c(z, pool_shape, top_down = None, theano_rng = None):
 
     hpart = []
     for i in xrange(r):
-        hpart.append( [ pt_ij / denom for pt_ij in pt[i] ] )
+        hpart.append([pt_ij / denom for pt_ij in pt[i]])
 
     h = T.alloc(0., batch_size, zr, zc, ch)
 
     for i in xrange(r):
         for j in xrange(c):
-            h = T.set_subtensor(h[:,i:zr:r,j:zc:c,:],hpart[i][j])
+            h = T.set_subtensor(h[:, i:zr:r, j:zc:c, :], hpart[i][j])
 
     h.name = 'h(%s)' % z_name
 
@@ -747,32 +770,37 @@ def max_pool_b01c(z, pool_shape, top_down = None, theano_rng = None):
                 events.append(hpart[i][j])
         events.append(off_prob)
 
-        events = [ event.dimshuffle(0,1,2,3,'x') for event in events ]
+        events = [event.dimshuffle(0, 1, 2, 3, 'x') for event in events]
 
         events = tuple(events)
 
-        stacked_events = T.concatenate( events, axis = 4)
+        stacked_events = T.concatenate(events, axis=4)
 
         batch_size, rows, cols, channels, outcomes = stacked_events.shape
-        reshaped_events = stacked_events.reshape((batch_size * rows * cols * channels, outcomes))
+        reshaped_events = stacked_events.reshape((batch_size * rows * cols *
+                                                  channels, outcomes))
 
-        multinomial = theano_rng.multinomial(pvals = reshaped_events, dtype = p.dtype)
+        multinomial = theano_rng.multinomial(pvals=reshaped_events,
+                                             dtype=p.dtype)
 
-        reshaped_multinomial = multinomial.reshape((batch_size, rows, cols, channels, outcomes))
+        reshaped_multinomial = multinomial.reshape((batch_size, rows, cols,
+                                                    channels, outcomes))
 
         h_sample = T.alloc(0., batch_size, zr, zc, ch)
 
         idx = 0
         for i in xrange(r):
             for j in xrange(c):
-                h_sample = T.set_subtensor(h_sample[:,i:zr:r,j:zc:c,:],
-                        reshaped_multinomial[:,:,:,:,idx])
+                h_sample = T.set_subtensor(h_sample[:, i:zr:r, j:zc:c, :],
+                                           reshaped_multinomial[:, :, :, :,
+                                           idx])
                 idx += 1
 
-        p_sample = 1 - reshaped_multinomial[:,:,:,:,-1]
+        p_sample = 1 - reshaped_multinomial[:, :, :, :, -1]
 
         return p, h, p_sample, h_sample
 
+
 def max_pool_softmax_with_bias_op(z, pool_shape):
     """
     .. todo::
@@ -784,7 +812,6 @@ def max_pool_softmax_with_bias_op(z, pool_shape):
     Also, this uses the ('b', 0, 1, 'c') format.
     """
 
-
     z_name = z.name
     if z_name is None:
         z_name = 'anon_z'
@@ -797,24 +824,25 @@ def max_pool_softmax_with_bias_op(z, pool_shape):
 
     for i in xrange(r):
         for j in xrange(c):
-            cur_part = z[:,i:zr:r,j:zc:c,:]
+            cur_part = z[:, i:zr:r, j:zc:c, :]
             assert cur_part.ndim == 4
             if z_name is not None:
-                cur_part.name = z_name + '[%d,%d]' % (i,j)
-            flat_z.append( cur_part.dimshuffle(0,1,2,3,'x') )
+                cur_part.name = z_name + '[%d,%d]' % (i, j)
+            flat_z.append(cur_part.dimshuffle(0, 1, 2, 3, 'x'))
 
     flat_z.append(T.zeros_like(flat_z[-1]))
 
-    stacked_z = T.concatenate( flat_z, axis = 4)
+    stacked_z = T.concatenate(flat_z, axis=4)
 
     batch_size, rows, cols, channels, outcomes = stacked_z.shape
-    reshaped_z = stacked_z.reshape((batch_size * rows * cols * channels, outcomes))
+    reshaped_z = stacked_z.reshape((batch_size * rows * cols * channels,
+                                    outcomes))
 
-    dist = T.nnet.softmax_with_bias(reshaped_z,T.zeros_like(reshaped_z[0,:]))
+    dist = T.nnet.softmax_with_bias(reshaped_z, T.zeros_like(reshaped_z[0, :]))
 
     dist = dist.reshape((batch_size, rows, cols, channels, outcomes))
 
-    p = 1. - dist[:,:,:,:,-1]
+    p = 1. - dist[:, :, :, :, -1]
     p.name = 'p(%s)' % z_name
 
     h = T.alloc(0., batch_size, zr, zc, ch)
@@ -822,14 +850,15 @@ def max_pool_softmax_with_bias_op(z, pool_shape):
     idx = 0
     for i in xrange(r):
         for j in xrange(c):
-            h = T.set_subtensor(h[:,i:zr:r,j:zc:c,:],
-                    dist[:,:,:,:,idx])
+            h = T.set_subtensor(h[:, i:zr:r, j:zc:c, :],
+                                dist[:, :, :, :, idx])
             idx += 1
 
     h.name = 'h(%s)' % z_name
 
     return p, h
 
+
 def max_pool_softmax_op(z, pool_shape):
     """
     .. todo::
@@ -853,24 +882,25 @@ def max_pool_softmax_op(z, pool_shape):
 
     for i in xrange(r):
         for j in xrange(c):
-            cur_part = z[:,i:zr:r,j:zc:c,:]
+            cur_part = z[:, i:zr:r, j:zc:c, :]
             assert cur_part.ndim == 4
             if z_name is not None:
-                cur_part.name = z_name + '[%d,%d]' % (i,j)
-            flat_z.append( cur_part.dimshuffle(0,1,2,3,'x') )
+                cur_part.name = z_name + '[%d,%d]' % (i, j)
+            flat_z.append(cur_part.dimshuffle(0, 1, 2, 3, 'x'))
 
     flat_z.append(T.zeros_like(flat_z[-1]))
 
-    stacked_z = T.concatenate( flat_z, axis = 4)
+    stacked_z = T.concatenate(flat_z, axis=4)
 
     batch_size, rows, cols, channels, outcomes = stacked_z.shape
-    reshaped_z = stacked_z.reshape((batch_size * rows * cols * channels, outcomes))
+    reshaped_z = stacked_z.reshape((batch_size * rows * cols * channels,
+                                    outcomes))
 
     dist = T.nnet.softmax(reshaped_z)
 
     dist = dist.reshape((batch_size, rows, cols, channels, outcomes))
 
-    p = 1. - dist[:,:,:,:,len(flat_z)-1]
+    p = 1. - dist[:, :, :, :, len(flat_z)-1]
     p.name = 'p(%s)' % z_name
 
     h = T.alloc(0., batch_size, zr, zc, ch)
@@ -878,8 +908,8 @@ def max_pool_softmax_op(z, pool_shape):
     idx = 0
     for i in xrange(r):
         for j in xrange(c):
-            h = T.set_subtensor(h[:,i:zr:r,j:zc:c,:],
-                    dist[:,:,:,:,idx])
+            h = T.set_subtensor(h[:, i:zr:r, j:zc:c, :],
+                                dist[:, :, :, :, idx])
             idx += 1
 
     h.name = 'h(%s)' % z_name
@@ -893,24 +923,25 @@ def profile(f):
 
         WRITEME
     """
-    print 'profiling ',f
-    rng = np.random.RandomState([2012,7,19])
+    print 'profiling ', f
+    rng = np.random.RandomState([2012, 7, 19])
     batch_size = 80
     rows = 26
     cols = 27
     channels = 30
     pool_rows = 2
     pool_cols = 3
-    zv = rng.randn( batch_size, rows, cols, channels ).astype(config.floatX)
+    zv = rng.randn(batch_size, rows, cols, channels).astype(config.floatX)
 
-    #put the inputs + outputs in shared variables so we don't pay GPU transfer during test
-    p_shared = sharedX(zv[:,0:rows:pool_rows,0:cols:pool_cols,:])
+    #put the inputs + outputs in shared variables so we don't pay GPU
+    # transfer during test
+    p_shared = sharedX(zv[:, 0:rows:pool_rows, 0:cols:pool_cols, :])
     h_shared = sharedX(zv)
     z_shared = sharedX(zv)
 
-    p_th, h_th = f( z_shared, (pool_rows, pool_cols) )
+    p_th, h_th = f(z_shared, (pool_rows, pool_cols))
 
-    func = function([],updates = { p_shared : p_th, h_shared : h_th} )
+    func = function([], updates={p_shared: p_th, h_shared: h_th})
 
     print 'warming up'
     for i in xrange(10):
@@ -926,7 +957,8 @@ def profile(f):
         t2 = time.time()
         print t2 - t1
         results.append(t2-t1)
-    print 'final: ',sum(results)/float(trials)
+    print 'final: ', sum(results)/float(trials)
+
 
 def profile_bc01(f):
     """
@@ -934,24 +966,25 @@ def profile_bc01(f):
 
         WRITEME
     """
-    print 'profiling ',f
-    rng = np.random.RandomState([2012,7,19])
+    print 'profiling ', f
+    rng = np.random.RandomState([2012, 7, 19])
     batch_size = 80
     rows = 26
     cols = 27
     channels = 30
     pool_rows = 2
     pool_cols = 3
-    zv = rng.randn( batch_size, channels, rows, cols).astype(config.floatX)
+    zv = rng.randn(batch_size, channels, rows, cols).astype(config.floatX)
 
-    #put the inputs + outputs in shared variables so we don't pay GPU transfer during test
-    p_shared = sharedX(zv[:,:,0:rows:pool_rows,0:cols:pool_cols])
+    # put the inputs + outputs in shared variables so we don't pay GPU
+    # transfer during test
+    p_shared = sharedX(zv[:, :, 0:rows:pool_rows, 0:cols:pool_cols])
     h_shared = sharedX(zv)
     z_shared = sharedX(zv)
 
-    p_th, h_th = f( z_shared, (pool_rows, pool_cols) )
+    p_th, h_th = f(z_shared, (pool_rows, pool_cols))
 
-    func = function([],updates = { p_shared : p_th, h_shared : h_th} )
+    func = function([], updates={p_shared: p_th, h_shared: h_th})
 
     print 'warming up'
     for i in xrange(10):
@@ -967,7 +1000,8 @@ def profile_bc01(f):
         t2 = time.time()
         print t2 - t1
         results.append(t2-t1)
-    print 'final: ',sum(results)/float(trials)
+    print 'final: ', sum(results)/float(trials)
+
 
 def profile_samples(f):
     """
@@ -975,8 +1009,8 @@ def profile_samples(f):
 
         WRITEME
     """
-    print 'profiling samples',f
-    rng = np.random.RandomState([2012,7,19])
+    print 'profiling samples', f
+    rng = np.random.RandomState([2012, 7, 19])
     theano_rng = MRG_RandomStreams(rng.randint(2147462579))
     batch_size = 80
     rows = 26
@@ -984,16 +1018,17 @@ def profile_samples(f):
     channels = 30
     pool_rows = 2
     pool_cols = 3
-    zv = rng.randn( batch_size, rows, cols, channels ).astype(config.floatX)
+    zv = rng.randn(batch_size, rows, cols, channels).astype(config.floatX)
 
-    #put the inputs + outputs in shared variables so we don't pay GPU transfer during test
-    p_shared = sharedX(zv[:,0:rows:pool_rows,0:cols:pool_cols,:])
+    #put the inputs + outputs in shared variables so we don't pay GPU
+    # transfer during test
+    p_shared = sharedX(zv[:, 0:rows:pool_rows, 0:cols:pool_cols, :])
     h_shared = sharedX(zv)
     z_shared = sharedX(zv)
 
-    p_th, h_th, ps_th, hs_th = f( z_shared, (pool_rows, pool_cols), theano_rng )
+    p_th, h_th, ps_th, hs_th = f(z_shared, (pool_rows, pool_cols), theano_rng)
 
-    func = function([],updates = { p_shared : ps_th, h_shared : hs_th} )
+    func = function([], updates={p_shared: ps_th, h_shared: hs_th})
 
     print 'warming up'
     for i in xrange(10):
@@ -1009,7 +1044,8 @@ def profile_samples(f):
         t2 = time.time()
         print t2 - t1
         results.append(t2-t1)
-    print 'final: ',sum(results)/float(trials)
+    print 'final: ', sum(results)/float(trials)
+
 
 def profile_grad(f):
     """
@@ -1017,23 +1053,25 @@ def profile_grad(f):
 
         WRITEME
     """
-    print 'profiling gradient of ',f
-    rng = np.random.RandomState([2012,7,19])
+    print 'profiling gradient of ', f
+    rng = np.random.RandomState([2012, 7, 19])
     batch_size = 80
     rows = 26
     cols = 27
     channels = 30
     pool_rows = 2
     pool_cols = 3
-    zv = rng.randn( batch_size, rows, cols, channels ).astype(config.floatX)
+    zv = rng.randn(batch_size, rows, cols, channels).astype(config.floatX)
 
-    #put the inputs + outputs in shared variables so we don't pay GPU transfer during test
+    #put the inputs + outputs in shared variables so we don't pay GPU
+    # transfer during test
     grad_shared = sharedX(zv)
     z_shared = sharedX(zv)
 
-    p_th, h_th = f( z_shared, (pool_rows, pool_cols) )
+    p_th, h_th = f(z_shared, (pool_rows, pool_cols))
 
-    func = function([],updates = { grad_shared : T.grad(p_th.sum() +  h_th.sum(), z_shared)} )
+    func = function([], updates={grad_shared: T.grad(p_th.sum() + h_th.sum(),
+                                 z_shared)})
 
     print 'warming up'
     for i in xrange(10):
@@ -1049,7 +1087,8 @@ def profile_grad(f):
         t2 = time.time()
         print t2 - t1
         results.append(t2-t1)
-    print 'final: ',sum(results)/float(trials)
+    print 'final: ', sum(results)/float(trials)
+
 
 def profile_grad_bc01(f):
     """
@@ -1057,23 +1096,25 @@ def profile_grad_bc01(f):
 
         WRITEME
     """
-    print 'profiling gradient of ',f
-    rng = np.random.RandomState([2012,7,19])
+    print 'profiling gradient of ', f
+    rng = np.random.RandomState([2012, 7, 19])
     batch_size = 80
     rows = 26
     cols = 27
     channels = 30
     pool_rows = 2
     pool_cols = 3
-    zv = rng.randn( batch_size, channels, rows, cols).astype(config.floatX)
+    zv = rng.randn(batch_size, channels, rows, cols).astype(config.floatX)
 
-    #put the inputs + outputs in shared variables so we don't pay GPU transfer during test
+    # put the inputs + outputs in shared variables so we don't pay GPU
+    # transfer during test
     grad_shared = sharedX(zv)
     z_shared = sharedX(zv)
 
-    p_th, h_th = f( z_shared, (pool_rows, pool_cols) )
+    p_th, h_th = f(z_shared, (pool_rows, pool_cols))
 
-    func = function([],updates = { grad_shared : T.grad(p_th.sum() +  h_th.sum(), z_shared)} )
+    func = function([], updates={grad_shared: T.grad(p_th.sum() + h_th.sum(),
+                                 z_shared)})
 
     print 'warming up'
     for i in xrange(10):
@@ -1089,7 +1130,8 @@ def profile_grad_bc01(f):
         t2 = time.time()
         print t2 - t1
         results.append(t2-t1)
-    print 'final: ',sum(results)/float(trials)
+    print 'final: ', sum(results)/float(trials)
+
 
 if __name__ == '__main__':
     profile_bc01(max_pool)
@@ -1104,8 +1146,3 @@ def profile_grad_bc01(f):
     profile_grad(max_pool_softmax_op)
     profile_grad(max_pool_softmax_with_bias_op)
     """
-
-
-
-
-


Test patch:
diff --git a/pylearn2/expr/tests/test_probabilistic_max_pooling.py b/pylearn2/expr/tests/test_probabilistic_max_pooling.py
--- a/pylearn2/expr/tests/test_probabilistic_max_pooling.py
+++ b/pylearn2/expr/tests/test_probabilistic_max_pooling.py
@@ -14,20 +14,22 @@
 from pylearn2.expr.probabilistic_max_pooling import max_pool_c01b
 from pylearn2.expr.probabilistic_max_pooling import max_pool_unstable
 from pylearn2.expr.probabilistic_max_pooling import max_pool_softmax_op
-from pylearn2.expr.probabilistic_max_pooling import max_pool_softmax_with_bias_op
+from pylearn2.expr.probabilistic_max_pooling import \
+    max_pool_softmax_with_bias_op
 from pylearn2.testing import no_debug_mode
 
 
 def check_correctness_channelwise(f):
+    """
+    Tests that the theano expression emitted by f computes the same values
+    as the ground truth python function
+    Note: to keep the python version as dead simple as possible (i.e., to make
+    sure there are not bugs in the ground truth) it uses the numerically
+    unstable verison of softmax. So this test does not work with too big of
+    numbers.
+    """
 
-    # Tests that the theano expression emitted by f computes the same values
-    # as the ground truth python function
-    # Note: to keep the python version as dead simple as possible (i.e., to make
-    # sure there are not bugs in the ground truth) it uses the numerically
-    # unstable verison of softmax. So this test does not work with too big of
-    # numbers.
-
-    rng = np.random.RandomState([2012,7,19])
+    rng = np.random.RandomState([2012, 7, 19])
     batch_size = 5
     pool_size = 4
     n = 3 * pool_size
@@ -50,29 +52,30 @@ def check_correctness_channelwise(f):
 
     assert p_np.shape == pv.shape
     assert h_np.shape == hv.shape
-    if not np.allclose(h_np,hv):
-        print (h_np.min(),h_np.max())
-        print (hv.min(),hv.max())
+    if not np.allclose(h_np, hv):
+        print (h_np.min(), h_np.max())
+        print (hv.min(), hv.max())
         assert False
-    if not np.allclose(p_np,pv):
+    if not np.allclose(p_np, pv):
         diff = abs(p_np - pv)
-        print 'max diff ',diff.max()
-        print 'min diff ',diff.min()
-        print 'ave diff ',diff.mean()
+        print 'max diff ', diff.max()
+        print 'min diff ', diff.min()
+        print 'ave diff ', diff.mean()
         assert False
 
-def check_correctness_sigmoid_channelwise(f):
 
-    # Tests that f is equivalent to the sigmoid function when the pool size is 1
+def check_correctness_sigmoid_channelwise(f):
+    """
+    Tests that f is equivalent to the sigmoid function when the pool size is 1
+    """
 
-    rng = np.random.RandomState([2012,7,19])
+    rng = np.random.RandomState([2012, 7, 19])
     batch_size = 5
     pool_size = 1
     n = 3 * pool_size
     zv = rng.randn(batch_size, n).astype(config.floatX) * 1. - 1.5
     top_down_v = rng.randn(batch_size,  n / pool_size).astype(config.floatX)
 
-
     z_th = T.matrix()
     z_th.name = 'z_th'
 
@@ -80,7 +83,7 @@ def check_correctness_sigmoid_channelwise(f):
     top_down_th.name = 'top_down_th'
 
     p_th, h_th = f(z_th, pool_size, top_down_th)
-    h_s = T.nnet.sigmoid(z_th +  top_down_th)
+    h_s = T.nnet.sigmoid(z_th + top_down_th)
 
     func = function([z_th, top_down_th], [p_th, h_th, h_s])
 
@@ -89,102 +92,114 @@ def check_correctness_sigmoid_channelwise(f):
 
     assert p_s.shape == pv.shape
     assert h_s.shape == hv.shape
-    if not np.allclose(h_s,hv):
-        print (h_s.min(),h_s.max())
-        print (hv.min(),hv.max())
+    if not np.allclose(h_s, hv):
+        print (h_s.min(), h_s.max())
+        print (hv.min(), hv.max())
         assert False
-    if not np.allclose(p_s,pv):
+    if not np.allclose(p_s, pv):
         diff = abs(p_s - pv)
-        print 'max diff ',diff.max()
-        print 'min diff ',diff.min()
-        print 'ave diff ',diff.mean()
+        print 'max diff ', diff.max()
+        print 'min diff ', diff.min()
+        print 'ave diff ', diff.mean()
         assert False
 
+
 def check_correctness(f):
-    rng = np.random.RandomState([2012,7,19])
+    rng = np.random.RandomState([2012, 7, 19])
     batch_size = 5
     rows = 32
     cols = 30
     channels = 3
     pool_rows = 2
     pool_cols = 3
-    zv = rng.randn( batch_size, rows, cols, channels ).astype(config.floatX) * 2. - 3.
+    zv = rng.randn(batch_size, rows, cols,
+                   channels).astype(config.floatX) * 2. - 3.
 
-    p_np, h_np = max_pool_python( zv, (pool_rows, pool_cols) )
+    p_np, h_np = max_pool_python(zv, (pool_rows, pool_cols))
 
-    z_th = T.TensorType( broadcastable=(False,False,False,False), dtype = config.floatX)()
+    z_th = T.TensorType(broadcastable=(False, False, False, False),
+                        dtype=config.floatX)()
     z_th.name = 'z_th'
 
-    p_th, h_th = f( z_th, (pool_rows, pool_cols) )
+    p_th, h_th = f(z_th, (pool_rows, pool_cols))
 
-    func = function([z_th],[p_th,h_th])
+    func = function([z_th], [p_th, h_th])
 
     pv, hv = func(zv)
 
     assert p_np.shape == pv.shape
     assert h_np.shape == hv.shape
-    if not np.allclose(h_np,hv):
-        print (h_np.min(),h_np.max())
-        print (hv.min(),hv.max())
+    if not np.allclose(h_np, hv):
+        print (h_np.min(), h_np.max())
+        print (hv.min(), hv.max())
         assert False
-    assert np.allclose(p_np,pv)
+    assert np.allclose(p_np, pv)
 
-def check_correctness_bc01(f):
 
-    # Tests that the theano expression emitted by f computes the same values
-    # as the ground truth python function
-    # Note: to keep the python version as dead simple as possible (i.e., to make
-    # sure there are not bugs in the ground truth) it uses the numerically
-    # unstable verison of softmax. So this test does not work with too big of
-    # numbers.
+def check_correctness_bc01(f):
+    """
+    Tests that the theano expression emitted by f computes the same values
+    as the ground truth python function
+    Note: to keep the python version as dead simple as possible (i.e., to make
+    sure there are not bugs in the ground truth) it uses the numerically
+    unstable verison of softmax. So this test does not work with too big of
+    numbers.
+    """
 
-    rng = np.random.RandomState([2012,7,19])
+    rng = np.random.RandomState([2012, 7, 19])
     batch_size = 5
     rows = 32
     cols = 30
     channels = 3
     pool_rows = 2
     pool_cols = 3
-    zv = rng.randn(batch_size,  rows, cols, channels).astype(config.floatX) * 1. - 1.5
-    top_down_v = rng.randn(batch_size,  rows / pool_rows, cols / pool_cols, channels).astype(config.floatX)
+    zv = rng.randn(batch_size,  rows, cols,
+                   channels).astype(config.floatX) * 1. - 1.5
+    top_down_v = rng.randn(batch_size, rows / pool_rows, cols / pool_cols,
+                           channels).astype(config.floatX)
 
     p_np, h_np = max_pool_python(zv, (pool_rows, pool_cols), top_down_v)
 
-    z_th = T.TensorType( broadcastable=(False,False,False,False), dtype = config.floatX)()
+    z_th = T.TensorType(broadcastable=(False, False, False, False),
+                        dtype = config.floatX)()
     z_th.name = 'z_th'
-    zr = z_th.dimshuffle(0,3,1,2)
+    zr = z_th.dimshuffle(0, 3, 1, 2)
 
-    top_down_th = T.TensorType( broadcastable=(False,False,False,False), dtype = config.floatX)()
+    top_down_th = T.TensorType(broadcastable=(False, False, False, False),
+                               dtype = config.floatX)()
     top_down_th.name = 'top_down_th'
-    top_down_r = top_down_th.dimshuffle(0,3,1,2)
+    top_down_r = top_down_th.dimshuffle(0, 3, 1, 2)
 
     p_th, h_th = f(zr, (pool_rows, pool_cols), top_down_r)
 
-    func = function([z_th, top_down_th], [p_th.dimshuffle(0,2,3,1), h_th.dimshuffle(0,2,3,1)])
+    func = function([z_th, top_down_th], [p_th.dimshuffle(0, 2, 3, 1),
+                                          h_th.dimshuffle(0, 2, 3, 1)])
 
     pv, hv = func(zv, top_down_v)
 
     assert p_np.shape == pv.shape
     assert h_np.shape == hv.shape
-    if not np.allclose(h_np,hv):
-        print (h_np.min(),h_np.max())
-        print (hv.min(),hv.max())
+    if not np.allclose(h_np, hv):
+        print (h_np.min(), h_np.max())
+        print (hv.min(), hv.max())
         assert False
-    if not np.allclose(p_np,pv):
+    if not np.allclose(p_np, pv):
         diff = abs(p_np - pv)
-        print 'max diff ',diff.max()
-        print 'min diff ',diff.min()
-        print 'ave diff ',diff.mean()
+        print 'max diff ', diff.max()
+        print 'min diff ', diff.min()
+        print 'ave diff ', diff.mean()
         assert False
 
-def check_correctness_c01b(f):
 
-    # Tests that the theano expression emitted by f computes the same values
-    # as the ground truth python function
-    # Note: to keep the python version as dead simple as possible (i.e., to make
-    # sure there are not bugs in the ground truth) it uses the numerically
-    # unstable version of softmax. So this test does not work with too big of
-    # numbers.
+def check_correctness_c01b(f):
+    """
+    Tests that the theano expression emitted by f computes the same values
+    as the ground truth python function
+    Note: to keep the python version as dead simple as possible (i.e., to make
+    sure there are not bugs in the ground truth) it uses the numerically
+    unstable version of softmax. So this test does not work with too big of
+    numbers.
+    """
 
     rng = np.random.RandomState([2013, 5, 6])
     batch_size = 5
@@ -195,63 +210,74 @@ def check_correctness_c01b(f):
     pool_cols = 3
 
     # Do the python ground truth in b01c format
-    zv = rng.randn(batch_size,  rows, cols, channels).astype(config.floatX) * 1. - 1.5
-    top_down_v = rng.randn(batch_size,  rows / pool_rows, cols / pool_cols, channels).astype(config.floatX)
+    zv = rng.randn(batch_size,  rows, cols,
+                   channels).astype(config.floatX) * 1. - 1.5
+    top_down_v = rng.randn(batch_size, rows / pool_rows, cols / pool_cols,
+                           channels).astype(config.floatX)
 
     p_np, h_np = max_pool_python(zv, (pool_rows, pool_cols), top_down_v)
 
     # Dimshuffle the inputs into c01b for the theano implementation
-    z_th = T.TensorType( broadcastable=(False,False,False,False), dtype = config.floatX)()
+    z_th = T.TensorType(broadcastable=(False, False, False, False),
+                        dtype = config.floatX)()
     z_th.tag.test_value = zv
     z_th.name = 'z_th'
-    zr = z_th.dimshuffle(3,1,2,0)
+    zr = z_th.dimshuffle(3, 1, 2, 0)
 
-    top_down_th = T.TensorType( broadcastable=(False,False,False,False), dtype = config.floatX)()
+    top_down_th = T.TensorType(broadcastable=(False, False, False, False),
+                               dtype = config.floatX)()
     top_down_th.name = 'top_down_th'
     top_down_th.tag.test_value = top_down_v
-    top_down_r = top_down_th.dimshuffle(3,1,2,0)
+    top_down_r = top_down_th.dimshuffle(3, 1, 2, 0)
 
     p_th, h_th = f(zr, (pool_rows, pool_cols), top_down_r)
 
-    func = function([z_th, top_down_th], [p_th.dimshuffle(3,1,2,0), h_th.dimshuffle(3,1,2,0)])
+    func = function([z_th, top_down_th], [p_th.dimshuffle(3, 1, 2, 0),
+                                          h_th.dimshuffle(3, 1, 2, 0)])
 
     pv, hv = func(zv, top_down_v)
 
     if not p_np.shape == pv.shape:
         raise AssertionError(str((p_np.shape, pv.shape)))
     assert h_np.shape == hv.shape
-    if not np.allclose(h_np,hv):
-        print (h_np.min(),h_np.max())
-        print (hv.min(),hv.max())
+    if not np.allclose(h_np, hv):
+        print (h_np.min(), h_np.max())
+        print (hv.min(), hv.max())
         assert False
-    if not np.allclose(p_np,pv):
+    if not np.allclose(p_np, pv):
         diff = abs(p_np - pv)
-        print 'max diff ',diff.max()
-        print 'min diff ',diff.min()
-        print 'ave diff ',diff.mean()
+        print 'max diff ', diff.max()
+        print 'min diff ', diff.min()
+        print 'ave diff ', diff.mean()
         assert False
     warnings.warn("TODO: make sampling tests run on c01b format of pooling.")
 
+
 @no_debug_mode
-def check_sample_correctishness(f):
+def check_sample_correctishness_b01c(f):
     batch_size = 5
     rows = 32
     cols = 30
     channels = 3
     pool_rows = 2
     pool_cols = 3
-    rng = np.random.RandomState([2012,9,26])
-    zv = rng.randn( batch_size, rows, cols, channels ).astype(config.floatX) * 2. - 3.
-    top_down_v = rng.randn( batch_size, rows / pool_rows, cols / pool_cols, channels ).astype(config.floatX)
-
-    z_th = T.TensorType( broadcastable=(False,False,False,False), dtype = config.floatX)()
+    rng = np.random.RandomState([2012, 9, 26])
+    zv = rng.randn(batch_size, rows, cols,
+                   channels).astype(config.floatX) * 2. - 3.
+    top_down_v = rng.randn(batch_size, rows / pool_rows, cols / pool_cols,
+                           channels).astype(config.floatX)
+
+    z_th = T.TensorType(broadcastable=(False, False, False, False),
+                        dtype = config.floatX)()
     z_th.name = 'z_th'
 
-    top_down_th = T.TensorType( broadcastable=(False,False,False,False), dtype = config.floatX)()
+    top_down_th = T.TensorType(broadcastable=(False, False, False, False),
+                               dtype = config.floatX)()
     top_down_th.name = 'top_down_th'
 
     theano_rng = MRG_RandomStreams(rng.randint(2147462579))
-    p_th, h_th, p_sth, h_sth = f( z_th, (pool_rows, pool_cols), top_down_th, theano_rng )
+    p_th, h_th, p_sth, h_sth = f(z_th, (pool_rows, pool_cols), top_down_th,
+                                 theano_rng)
 
     prob_func = function([z_th, top_down_th], [p_th, h_th])
     pv, hv = prob_func(zv, top_down_v)
@@ -261,20 +287,19 @@ def check_sample_correctishness(f):
     acc_p = 0. * pv
     acc_h = 0. * hv
 
-    # make sure the test gets good coverage, ie, that it includes many different
-    # activation probs for both detector and pooling layer
+    # make sure the test gets good coverage, ie, that it includes many
+    # different activation probs for both detector and pooling layer
     buckets = 10
     bucket_width = 1. / float(buckets)
     for i in xrange(buckets):
         lower_lim = i * bucket_width
         upper_lim = (i+1) * bucket_width
 
-        assert np.any( (pv >= lower_lim) * (pv < upper_lim) )
-        assert np.any( (hv >= lower_lim) * (hv < upper_lim) )
+        assert np.any((pv >= lower_lim) * (pv < upper_lim))
+        assert np.any((hv >= lower_lim) * (hv < upper_lim))
 
     assert upper_lim == 1.
 
-
     for i in xrange(10000):
         ps, hs = sample_func(zv, top_down_v)
 
@@ -291,17 +316,20 @@ def check_sample_correctishness(f):
     hd = np.abs(est_h-hv)
 
     """
-    # plot maps of the estimation error, this is to see if it has some spatial pattern
-    # this is useful for detecting bugs like not handling the border correctly, etc.
+    # plot maps of the estimation error, this is to see if it has
+    # some spatial pattern this is useful for detecting bugs like
+    # not handling the border correctly, etc.
     from pylearn2.gui.patch_viewer import PatchViewer
 
-    pv = PatchViewer((pd.shape[0],pd.shape[3]),(pd.shape[1],pd.shape[2]),is_color = False)
+    pv = PatchViewer((pd.shape[0],pd.shape[3]),(pd.shape[1],pd.shape[2]),
+                      is_color = False)
     for i in xrange(pd.shape[0]):
     for j in xrange(pd.shape[3]):
-    pv.add_patch( (pd[i,:,:,j] / pd.max() )* 2.0 - 1.0, rescale = False)
+    pv.add_patch((pd[i,:,:,j] / pd.max() )* 2.0 - 1.0, rescale = False)
     pv.show()
 
-    pv = PatchViewer((hd.shape[0],hd.shape[3]),(hd.shape[1],hd.shape[2]),is_color = False)
+    pv = PatchViewer((hd.shape[0],hd.shape[3]),(hd.shape[1],hd.shape[2]),
+                      is_color = False)
     for i in xrange(hd.shape[0]):
     for j in xrange(hd.shape[3]):
     pv.add_patch( (hd[i,:,:,j] / hd.max() )* 2.0 - 1.0, rescale = False)
@@ -331,29 +359,120 @@ def check_sample_correctishness(f):
     assert max(pd.max(), hd.max()) < .17
 
     # Do exhaustive checks on just the last sample
-    assert np.all( (ps ==0) + (ps == 1) )
-    assert np.all( (hs == 0) + (hs == 1) )
+    assert np.all((ps == 0) + (ps == 1))
+    assert np.all((hs == 0) + (hs == 1))
 
     for k in xrange(batch_size):
         for i in xrange(ps.shape[1]):
             for j in xrange(ps.shape[2]):
                 for l in xrange(channels):
-                    p = ps[k,i,j,l]
-                    h = hs[k,i*pool_rows:(i+1)*pool_rows,j*pool_cols:(j+1)*pool_cols,l]
+                    p = ps[k, i, j, l]
+                    h = hs[k, i*pool_rows:(i+1)*pool_rows,
+                           j*pool_cols:(j+1)*pool_cols, l]
                     assert h.shape == (pool_rows, pool_cols)
                     assert p == h.max()
 
+    """ If you made it to here, it's correctish
+     (cant tell if samples are perfectly "correct") """
+
+
+@no_debug_mode
+def check_sample_correctishness_c01b(f):
+    batch_size = 5
+    rows = 32
+    cols = 30
+    channels = 3
+    pool_rows = 2
+    pool_cols = 3
+    rng = np.random.RandomState([2012, 9, 26])
+    zv = rng.randn(channels, rows, cols,
+                   batch_size).astype(config.floatX) * 2. - 3.
+    top_down_v = rng.randn(channels, rows / pool_rows, cols / pool_cols,
+                           batch_size).astype(config.floatX)
+
+    z_th = T.TensorType(broadcastable=(False, False, False, False),
+                        dtype = config.floatX)()
+    z_th.name = 'z_th'
+    z_th.tag.test_value = zv
+
+    top_down_th = T.TensorType(broadcastable=(False, False, False, False),
+                               dtype = config.floatX)()
+    top_down_th.name = 'top_down_th'
+    top_down_th.tag.test_value = top_down_v
+
+    theano_rng = MRG_RandomStreams(rng.randint(2147462579))
+    p_th, h_th, p_sth, h_sth = f(z_th, (pool_rows, pool_cols), top_down_th,
+                                 theano_rng)
+
+    prob_func = function([z_th, top_down_th], [p_th, h_th])
+    pv, hv = prob_func(zv, top_down_v)
+
+    sample_func = function([z_th, top_down_th], [p_sth, h_sth])
+
+    acc_p = 0. * pv
+    acc_h = 0. * hv
+
+    # make sure the test gets good coverage, ie, that it includes
+    # many different activation probs for both detector and pooling layer
+    buckets = 10
+    bucket_width = 1. / float(buckets)
+    for i in xrange(buckets):
+        lower_lim = i * bucket_width
+        upper_lim = (i+1) * bucket_width
+
+        assert np.any((pv >= lower_lim) * (pv < upper_lim))
+        assert np.any((hv >= lower_lim) * (hv < upper_lim))
+
+    assert upper_lim == 1.
+
+    for i in xrange(10000):
+        ps, hs = sample_func(zv, top_down_v)
+
+        assert ps.shape == pv.shape
+        assert hs.shape == hv.shape
+
+        acc_p += ps
+        acc_h += hs
+
+    est_p = acc_p / float(i+1)
+    est_h = acc_h / float(i+1)
+
+    pd = np.abs(est_p-pv)
+    hd = np.abs(est_h-hv)
+
+    # don't really know how tight this should be
+    # but you can try to pose an equivalent problem
+    # and implement it in another way
+    # using a numpy implementation in softmax_acc.py
+    # I got a max error of .17
+    assert max(pd.max(), hd.max()) < .17
+
+    # Do exhaustive checks on just the last sample
+    assert np.all((ps == 0) + (ps == 1))
+    assert np.all((hs == 0) + (hs == 1))
+
+    for k in xrange(batch_size):
+        for i in xrange(ps.shape[1]):
+            for j in xrange(ps.shape[2]):
+                for l in xrange(channels):
+                    p = ps[l, i, j, k]
+                    h = hs[l, i*pool_rows:(i+1)*pool_rows,
+                           j*pool_cols:(j+1)*pool_cols, k]
+                    assert h.shape == (pool_rows, pool_cols)
+                    assert p == h.max()
 
     """ If you made it to here, it's correctish
      (cant tell if samples are perfectly "correct") """
 
+
 @no_debug_mode
 def check_sample_correctishness_bc01(f):
-
-    # Tests that the sample mean converges to the conditional expectation given by the
-    # function
-    # Tests that p really is the max of the samples
-    # Tests that at most one h in a group is on
+    """
+    Tests that the sample mean converges to the conditional
+    expectation given by the function
+    Tests that p really is the max of the samples
+    Tests that at most one h in a group is on
+    """
 
     batch_size = 5
     rows = 32
@@ -362,20 +481,25 @@ def check_sample_correctishness_bc01(f):
     pool_rows = 2
     pool_cols = 3
 
-    rng = np.random.RandomState([2012,9,26])
-    zv = rng.randn( batch_size, channels, rows, cols).astype(config.floatX) * 2. - 3.
-    top_down_v = rng.randn( batch_size, channels, rows / pool_rows, cols / pool_cols).astype(config.floatX)
+    rng = np.random.RandomState([2012, 9, 26])
+    zv = rng.randn(batch_size, channels, rows,
+                   cols).astype(config.floatX) * 2. - 3.
+    top_down_v = rng.randn(batch_size, channels, rows / pool_rows,
+                           cols / pool_cols).astype(config.floatX)
 
-    z_th = T.TensorType( broadcastable=(False,False,False,False), dtype = config.floatX)()
+    z_th = T.TensorType(broadcastable=(False, False, False, False),
+                        dtype = config.floatX)()
     z_th.tag.test_value = zv
     z_th.name = 'z_th'
 
-    top_down_th = T.TensorType( broadcastable=(False,False,False,False), dtype = config.floatX)()
+    top_down_th = T.TensorType(broadcastable=(False, False, False, False),
+                               dtype = config.floatX)()
     top_down_th.tag.test_value = top_down_v
     top_down_th.name = 'top_down_th'
 
     theano_rng = MRG_RandomStreams(rng.randint(2147462579))
-    p_th, h_th, p_sth, h_sth = f( z_th, (pool_rows, pool_cols), top_down_th, theano_rng )
+    p_th, h_th, p_sth, h_sth = f(z_th, (pool_rows, pool_cols), top_down_th,
+                                 theano_rng)
 
     prob_func = function([z_th, top_down_th], [p_th, h_th])
     pv, hv = prob_func(zv, top_down_v)
@@ -385,20 +509,19 @@ def check_sample_correctishness_bc01(f):
     acc_p = 0. * pv
     acc_h = 0. * hv
 
-    # make sure the test gets good coverage, ie, that it includes many different
-    # activation probs for both detector and pooling layer
+    # make sure the test gets good coverage, ie, that it includes many
+    # different activation probs for both detector and pooling layer
     buckets = 10
     bucket_width = 1. / float(buckets)
     for i in xrange(buckets):
         lower_lim = i * bucket_width
         upper_lim = (i+1) * bucket_width
 
-        assert np.any( (pv >= lower_lim) * (pv < upper_lim) )
-        assert np.any( (hv >= lower_lim) * (hv < upper_lim) )
+        assert np.any((pv >= lower_lim) * (pv < upper_lim))
+        assert np.any((hv >= lower_lim) * (hv < upper_lim))
 
     assert upper_lim == 1.
 
-
     for i in xrange(10000):
         ps, hs = sample_func(zv, top_down_v)
 
@@ -415,17 +538,20 @@ def check_sample_correctishness_bc01(f):
     hd = np.abs(est_h-hv)
 
     """
-    # plot maps of the estimation error, this is to see if it has some spatial pattern
-    # this is useful for detecting bugs like not handling the border correctly, etc.
+    # plot maps of the estimation error, this is to see if it has some
+    # spatial pattern this is useful for detecting bugs like not handling
+    # the border correctly, etc.
     from pylearn2.gui.patch_viewer import PatchViewer
 
-    pv = PatchViewer((pd.shape[0],pd.shape[3]),(pd.shape[1],pd.shape[2]),is_color = False)
+    pv = PatchViewer((pd.shape[0],pd.shape[3]),(pd.shape[1],pd.shape[2]),
+                                                            is_color = False)
     for i in xrange(pd.shape[0]):
     for j in xrange(pd.shape[3]):
     pv.add_patch( (pd[i,:,:,j] / pd.max() )* 2.0 - 1.0, rescale = False)
     pv.show()
 
-    pv = PatchViewer((hd.shape[0],hd.shape[3]),(hd.shape[1],hd.shape[2]),is_color = False)
+    pv = PatchViewer((hd.shape[0],hd.shape[3]), (hd.shape[1],hd.shape[2]),
+                                                        is_color = False)
     for i in xrange(hd.shape[0]):
     for j in xrange(hd.shape[3]):
     pv.add_patch( (hd[i,:,:,j] / hd.max() )* 2.0 - 1.0, rescale = False)
@@ -455,38 +581,39 @@ def check_sample_correctishness_bc01(f):
     assert max(pd.max(), hd.max()) < .17
 
     # Do exhaustive checks on just the last sample
-    assert np.all( (ps ==0) + (ps == 1) )
-    assert np.all( (hs == 0) + (hs == 1) )
+    assert np.all((ps == 0) + (ps == 1))
+    assert np.all((hs == 0) + (hs == 1))
 
     for k in xrange(batch_size):
         for i in xrange(ps.shape[2]):
             for j in xrange(ps.shape[3]):
                 for l in xrange(channels):
-                    p = ps[k,l,i,j]
-                    h = hs[k,l,i*pool_rows:(i+1)*pool_rows,j*pool_cols:(j+1)*pool_cols]
+                    p = ps[k, l, i, j]
+                    h = hs[k, l, i*pool_rows:(i+1)*pool_rows,
+                           j*pool_cols:(j+1)*pool_cols]
                     assert h.shape == (pool_rows, pool_cols)
                     assert p == h.max()
                     assert h.sum() <= 1
 
-
     """ If you made it to here, it's correctish
      (cant tell if samples are perfectly "correct") """
 
+
 @no_debug_mode
 def check_sample_correctishness_channelwise(f):
-
-    # Tests that the sample mean converges to the conditional expectation given by the
-    # function
-    # Tests that p really is the max of the samples
-    # Tests that at most one h in a group is on
+    """
+    Tests that the sample mean converges to the conditional expectation given
+    by the function Tests that p really is the max of the samples tests that
+    at most one h in a group is on
+    """
 
     batch_size = 27
     pool_size = 4
     n = pool_size * 21
 
-    rng = np.random.RandomState([2012,9,26])
-    zv = rng.randn( batch_size, n).astype(config.floatX) * 3.5 - 5.
-    top_down_v = rng.randn( batch_size, n / pool_size).astype(config.floatX)
+    rng = np.random.RandomState([2012, 9, 26])
+    zv = rng.randn(batch_size, n).astype(config.floatX) * 3.5 - 5.
+    top_down_v = rng.randn(batch_size, n / pool_size).astype(config.floatX)
 
     z_th = T.matrix()
     z_th.tag.test_value = zv
@@ -507,8 +634,8 @@ def check_sample_correctishness_channelwise(f):
     acc_p = 0. * pv
     acc_h = 0. * hv
 
-    # make sure the test gets good coverage, ie, that it includes many different
-    # activation probs for both detector and pooling layer
+    # make sure the test gets good coverage, ie, that it includes
+    # many different activation probs for both detector and pooling layer
     buckets = 10
     bucket_width = 1. / float(buckets)
     print pv.min(), pv.max()
@@ -518,12 +645,11 @@ def check_sample_correctishness_channelwise(f):
         upper_lim = (i+1) * bucket_width
         print lower_lim, upper_lim
 
-        assert np.any( (pv >= lower_lim) * (pv < upper_lim) )
-        assert np.any( (hv >= lower_lim) * (hv < upper_lim) )
+        assert np.any((pv >= lower_lim) * (pv < upper_lim))
+        assert np.any((hv >= lower_lim) * (hv < upper_lim))
 
     assert upper_lim == 1.
 
-
     for i in xrange(10000):
         ps, hs = sample_func(zv, top_down_v)
 
@@ -540,17 +666,20 @@ def check_sample_correctishness_channelwise(f):
     hd = np.abs(est_h-hv)
 
     """
-    # plot maps of the estimation error, this is to see if it has some spatial pattern
-    # this is useful for detecting bugs like not handling the border correctly, etc.
-    from pylearn2.gui.patch_viewer import PatchViewer
+    # plot maps of the estimation error, this is to see if it has some
+    # spatial pattern this is useful for detecting bugs like not handling
+    # the border correctly, etc.
+    # from pylearn2.gui.patch_viewer import PatchViewer
 
-    pv = PatchViewer((pd.shape[0],pd.shape[3]),(pd.shape[1],pd.shape[2]),is_color = False)
+    pv = PatchViewer((pd.shape[0],pd.shape[3]),(pd.shape[1],pd.shape[2]),
+                                                            is_color = False)
     for i in xrange(pd.shape[0]):
     for j in xrange(pd.shape[3]):
     pv.add_patch( (pd[i,:,:,j] / pd.max() )* 2.0 - 1.0, rescale = False)
     pv.show()
 
-    pv = PatchViewer((hd.shape[0],hd.shape[3]),(hd.shape[1],hd.shape[2]),is_color = False)
+    pv = PatchViewer((hd.shape[0],hd.shape[3]),(hd.shape[1],hd.shape[2]),
+                                                            is_color = False)
     for i in xrange(hd.shape[0]):
     for j in xrange(hd.shape[3]):
     pv.add_patch( (hd[i,:,:,j] / hd.max() )* 2.0 - 1.0, rescale = False)
@@ -580,51 +709,64 @@ def check_sample_correctishness_channelwise(f):
     assert max(pd.max(), hd.max()) < .17
 
     # Do exhaustive checks on just the last sample
-    assert np.all( (ps ==0) + (ps == 1) )
-    assert np.all( (hs == 0) + (hs == 1) )
+    assert np.all((ps == 0) + (ps == 1))
+    assert np.all((hs == 0) + (hs == 1))
 
     for k in xrange(batch_size):
         for i in xrange(ps.shape[1]):
-            p = ps[k,i]
-            h = hs[k,i*pool_size:(i+1)*pool_size]
+            p = ps[k, i]
+            h = hs[k, i*pool_size:(i+1)*pool_size]
             assert h.shape == (pool_size,)
             assert p == h.max()
             assert h.sum() <= 1
 
-
     """ If you made it to here, it's correctish
      (cant tell if samples are perfectly "correct") """
 
+
 def test_max_pool_channels():
     check_correctness_channelwise(max_pool_channels)
 
+
 def test_max_pool_channels_sigmoid():
     check_correctness_sigmoid_channelwise(max_pool_channels)
 
+
 def test_max_pool_channels_samples():
     check_sample_correctishness_channelwise(max_pool_channels)
 
+
 def test_max_pool():
     check_correctness_bc01(max_pool)
 
+
 def test_max_pool_c01b():
     check_correctness_c01b(max_pool_c01b)
 
+
 def test_max_pool_samples():
     check_sample_correctishness_bc01(max_pool)
 
+
 def test_max_pool_b01c_samples():
-    check_sample_correctishness(max_pool_b01c)
+    check_sample_correctishness_b01c(max_pool_b01c)
+
+
+def test_max_pool_c01b_samples():
+    check_sample_correctishness_c01b(max_pool_c01b)
+
 
 def test_max_pool_b01c():
     check_correctness(max_pool_b01c)
 
+
 def test_max_pool_unstable():
     check_correctness(max_pool_unstable)
 
+
 def test_max_pool_softmax_op():
     check_correctness(max_pool_softmax_op)
 
+
 def test_max_pool_softmax_with_bias_op():
     check_correctness(max_pool_softmax_with_bias_op)
-