Merge pull request #4154 from abergeron/fix_blocksparse

theano/sandbox/cuda/blocksparse.py

@@ -23,21 +23,13 @@ class GpuSparseBlockGemv(GpuOp):
     to change without notice.  Use the sandbox.blocksparse.sparse_block_dot()
     function for a stable interface.
     """
+    __props__ = ('inplace',)
 
     def __init__(self, inplace=False):
         self.inplace = inplace
         if self.inplace:
             self.destroy_map = {0: [0]}
 
-    def __eq__(self, other):
-        return type(self) == type(other) and self.inplace == other.inplace
-
-    def __hash__(self):
-        return hash(type(self)) ^ hash(self.inplace)
-
-    def __str__(self):
-        return "GpuSparseBlockGemv%s" % ("{inplace}" if self.inplace else "")
-
     def make_node(self, o, W, h, inputIdx, outputIdx):
         o = basic_ops.as_cuda_ndarray_variable(o)
         W = basic_ops.as_cuda_ndarray_variable(W)
@@ -350,28 +342,20 @@ gpu_sparse_block_gemv_inplace = GpuSparseBlockGemv(True)
 
 class GpuSparseBlockOuter(GpuOp):
     """
-    CPU version of SparseBlockOuter. See SparseBlockOuter's docstring for more
+    GPU version of SparseBlockOuter. See SparseBlockOuter's docstring for more
     information.
 
     This op should not be called directly since its interface is
     subject to change without notice.  It is involved in the gradient
     of GpuSparseBlockGemv. The gradient is not implemented.
     """
+    __props__ = ('inplace',)
 
     def __init__(self, inplace=False):
         self.inplace = inplace
         if self.inplace:
             self.destroy_map = {0: [0]}
 
-    def __eq__(self, other):
-        return type(self) == type(other) and self.inplace == other.inplace
-
-    def __hash__(self):
-        return hash(type(self)) ^ hash(self.inplace)
-
-    def __str__(self):
-        return "GpuSparseBlockOuter%s" % ("{inplace}" if self.inplace else "")
-
     def make_node(self, o, x, y, xIdx, yIdx, alpha=None):
         one = tensor.constant(numpy.asarray(1.0, dtype='float32'))
         o = basic_ops.as_cuda_ndarray_variable(o)

theano/sandbox/cuda/tests/test_blocksparse.py

@@ -7,7 +7,8 @@ import theano.tests.unittest_tools as utt
 import theano.tensor.nnet.tests.test_blocksparse
 
 import theano.sandbox.cuda as cuda_ndarray
-from theano.sandbox.cuda.blocksparse import (GpuSparseBlockOuter,
+from theano.sandbox.cuda.blocksparse import (GpuSparseBlockGemv,
+                                             GpuSparseBlockOuter,
                                              gpu_sparse_block_gemv,
                                              gpu_sparse_block_outer)
 from theano.sandbox.cuda.var import float32_shared_constructor
@@ -28,6 +29,8 @@ class BlockSparse_Gemv_and_Outer(
         self.mode = mode_with_gpu.excluding('constant_folding')
         self.gemv_op = gpu_sparse_block_gemv
         self.outer_op = gpu_sparse_block_outer
+        self.gemv_class = GpuSparseBlockGemv
+        self.outer_class = GpuSparseBlockOuter
 
     # This test is temporarily disabled since we disabled the output_merge
     # and alpha_merge optimizations for blocksparse due to brokeness.

theano/tensor/nnet/blocksparse.py

@@ -22,6 +22,7 @@ class SparseBlockGemv(Op):
         :scale: 50 %
 
     """
+    __props__ = ('inplace',)
 
     registered_opts = []
 
@@ -90,10 +91,7 @@ class SparseBlockGemv(Op):
         assert inputIdx.type.dtype in discrete_dtypes
         assert outputIdx.type.dtype in discrete_dtypes
 
-        output = o.type.__class__(dtype=o.type.dtype,
-                                  broadcastable=(False,) * o.ndim)()
-
-        return Apply(self, [o, W, h, inputIdx, outputIdx], [output])
+        return Apply(self, [o, W, h, inputIdx, outputIdx], [o.type()])
 
     def perform(self, node, inp, out_):
         o, W, h, iIdx, oIdx = inp[:5]
@@ -110,6 +108,9 @@ class SparseBlockGemv(Op):
                     o[b, j, :] += numpy.dot(h[b, i], w)
         out_[0][0] = o
 
+    def infer_shape(self, node, input_shapes):
+        return [input_shapes[0]]
+
     def grad(self, inputs, grads):
         o, W, h, inputIdx, outputIdx = inputs
         go = grads[0]
@@ -138,6 +139,7 @@ class SparseBlockOuter(Op):
     This op is involved in the gradient of SparseBlockGemv.
 
     """
+    __props__ = ('inplace',)
 
     registered_opts = []
 
@@ -190,11 +192,11 @@ class SparseBlockOuter(Op):
         if alpha is None:
             alpha = one
 
-        output = o.type.__class__(dtype=o.type.dtype,
-                                  broadcastable=(False,) * o.ndim)()
-
         return Apply(self, [o, x, y, xIdx, yIdx, alpha],
-                     [output])
+                     [o.type()])
+
+    def infer_shape(self, node, input_shapes):
+        return [input_shapes[0]]
 
     def perform(self, node, inp, out_):
         o, x, y, xIdx, yIdx, alpha = inp[:6]

theano/tensor/nnet/nnet.py

@@ -2302,7 +2302,7 @@ def h_softmax(x, batch_size, n_outputs, n_classes, n_outputs_per_class,
         output_probs = theano.tensor.nnet.softmax(
             activations.reshape((-1, n_outputs_per_class)))
         output_probs = output_probs.reshape((batch_size, n_classes, -1))
-        output_probs = class_probs[:, :, None] * output_probs
+        output_probs = class_probs.dimshuffle(0, 1, 'x') * output_probs
         output_probs = output_probs.reshape((batch_size, -1))
         # output_probs.shape[1] is n_classes * n_outputs_per_class, which might
         # be greater than n_outputs, so we ignore the potential irrelevant
@@ -2321,11 +2321,11 @@ def h_softmax(x, batch_size, n_outputs, n_classes, n_outputs_per_class,
 
         # Second softmax that computes the output probabilities
         activations = sparse_block_dot(
-            W2[None, :, :, :], x[:, None, :],
+            W2.dimshuffle('x', 0, 1, 2), x.dimshuffle(0, 'x', 1),
             tensor.zeros((batch_size, 1), dtype='int32'), b2,
-            target_classes[:, None])
+            target_classes.dimshuffle(0, 'x'))
 
-        output_probs = theano.tensor.nnet.softmax(activations[:, 0, :])
+        output_probs = theano.tensor.nnet.softmax(activations.dimshuffle(0, 2))
         target_class_probs = class_probs[tensor.arange(batch_size),
                                          target_classes]
         output_probs = output_probs[tensor.arange(batch_size),

theano/tensor/nnet/tests/test_blocksparse.py

 """
     Tests for block sparse dot
 """
-import unittest
-
 import numpy
 from numpy.random import randn
 
@@ -10,15 +8,12 @@ import theano
 from theano import tensor
 import theano.tests.unittest_tools as utt
 
-from theano.tensor.nnet.blocksparse import sparse_block_dot, \
-    sparse_block_gemv, sparse_block_outer
-
-
-class BlockSparse_Gemv_and_Outer(unittest.TestCase):
+from theano.tensor.nnet.blocksparse import (
+    sparse_block_dot, sparse_block_gemv, sparse_block_outer,
+    SparseBlockGemv, SparseBlockOuter)
 
-    def runTest(self):
-        pass
 
+class BlockSparse_Gemv_and_Outer(utt.InferShapeTester):
     def setUp(self):
         utt.seed_rng()
         mode = None
@@ -29,6 +24,8 @@ class BlockSparse_Gemv_and_Outer(unittest.TestCase):
         )
         self.gemv_op = sparse_block_gemv
         self.outer_op = sparse_block_outer
+        self.gemv_class = SparseBlockGemv
+        self.outer_class = SparseBlockOuter
 
     @staticmethod
     def gemv_data():
@@ -280,3 +277,40 @@ class BlockSparse_Gemv_and_Outer(unittest.TestCase):
             o_val, x_val, y_val, xIdx_val, yIdx_val)
 
         utt.assert_allclose(ref_out, th_out)
+
+    def test_dot_infershape(self):
+        b = tensor.fmatrix()
+        W = tensor.ftensor4()
+        h = tensor.ftensor3()
+        iIdx = tensor.imatrix()
+        oIdx = tensor.imatrix()
+
+        self._compile_and_check([W, h, iIdx, b, oIdx],
+                                [sparse_block_dot(W, h, iIdx, b, oIdx)],
+                                self.gemv_data(),
+                                self.gemv_class)
+
+    def test_gemv_infershape(self):
+        b = tensor.fmatrix()
+        W = tensor.ftensor4()
+        h = tensor.ftensor3()
+        iIdx = tensor.imatrix()
+        oIdx = tensor.imatrix()
+
+        self._compile_and_check(
+            [W, h, iIdx, b, oIdx],
+            [self.gemv_op(b.take(oIdx, axis=0), W, h, iIdx, oIdx)],
+            self.gemv_data(),
+            self.gemv_class)
+
+    def test_outer_infershape(self):
+        o = tensor.ftensor4()
+        x = tensor.ftensor3()
+        y = tensor.ftensor3()
+        xIdx = tensor.imatrix()
+        yIdx = tensor.imatrix()
+
+        self._compile_and_check([o, x, y, xIdx, yIdx],
+                                [self.outer_op(o, x, y, xIdx, yIdx)],
+                                self.outer_data(),
+                                self.outer_class)