Merge pull request #2965 from briancheung/batched_dot_grad

theano/sandbox/cuda/blas.py

@@ -199,6 +199,20 @@ class BatchedDotOp(GpuOp):
             } while (0)
         """
 
+    def grad(self, inp, grads):
+        x, y = inp
+        gz, = grads
+
+        xgrad = batched_dot(gz, y.dimshuffle(0, 2, 1))
+        ygrad = batched_dot(x.dimshuffle(0, 2, 1), gz)
+
+        rval = xgrad, ygrad
+
+        for elem in rval:
+            assert elem.dtype.find('float') != -1
+
+        return rval
+
 batched_dot = BatchedDotOp()
 
 class GpuDot22(GpuOp):

theano/sandbox/cuda/tests/test_blas.py

@@ -43,6 +43,7 @@ mode_without_gpu.check_py_code = False
 def my_rand(*shape):
     return theano._asarray(numpy.random.rand(*shape), dtype='float32')
 
+
 class TestBatchedDot(TestCase):
 
     def test_batched_dot_correctness(self):
@@ -106,6 +107,14 @@ class TestBatchedDot(TestCase):
         # Shape mismatch
         self.assertRaises(RuntimeError, fail, (5,4,3), (5,2,2))
 
+    def test_batched_dot_gradient(self):
+        theano.tests.unittest_tools.verify_grad(
+            batched_dot,
+            [numpy.random.randn(5,7,2).astype(numpy.float32),
+             numpy.random.randn(5,2,6).astype(numpy.float32)],
+            mode=mode_with_gpu)
+
+
 def test_dot22():
     def cmp(a_shp, b_shp):
         a0 = my_rand(*a_shp)