Repository: Theano/Theano

Issue:
batched_dot fails to compile with the new gpuarray backend when the second tensor doesn't have at least 3 dimensions
For example, the following runs without errors when `device=gpu`. However, the following prints an optimization error when `device=cuda`. The program doesn't break though since the error is caught and a non-gpuarray backend is probably used. I'm using the release candidate `rel-0.9.0rc4`.
```
import theano
import theano.tensor as T
A_var = T.tensor3('A')
b_var = T.matrix('b')
c_var = T.batched_dot(A_var, b_var)
c_fn = theano.function([A_var, b_var], c_var)  # optimization error is printed when THEANO_FLAGS=device=cuda
d_var = T.batched_dot(b_var, b_var)
d_fn = theano.function([b_var], d_var)  # optimization error is printed when THEANO_FLAGS=device=cuda
```
The issues is caused because batched_dot ends up using [`local_gpua_gemmbatch`](https://github.com/Theano/Theano/blob/master/theano/gpuarray/opt.py#L1186-L1189), which assumes the input `b` to have at least 3 dimensions (since `b.shape[2]` is called).


Discussion:
Indeed, we should probably do the same trick as [on the CPU](https://github.com/Theano/Theano/blob/52903f8267cff316fc669e207eac4e2ecae952a6/theano/tensor/blas.py#L2221), either in the Op itself or in the graph during the optimization.
Even making the tensor 3D outside of the Op is not enough:
```
RuntimeError: cublasSgemmBatched: (cublas) Invalid value.
Apply node that caused the error: GpuGemmBatch{inplace=False}(GpuAllocEmpty{dtype='float32', context_name=None}.0, TensorConstant{1.0}, hyp_bilstm_apply_states_copy[Gpua], InplaceGpuDimShuffle{0,1,x}.0, TensorConstant{0.0})
Toposort index: 8
Inputs types: [GpuArrayType<None>(float32, (False, False, True)), TensorType(float32, scalar), GpuArrayType<None>(float32, 3D), GpuArrayType<None>(float32, (False, False, True)), TensorType(float32, scalar)]
Inputs shapes: [(32, 19, 1), (), (32, 19, 600), (32, 600, 1), ()]
Inputs strides: [(76, 4, 4), (), (45600, 2400, 4), (98400, 4, 4), ()]
```
```python
# problem.py
import numpy as np
from theano import tensor, function, config

x = tensor.tensor3('x')
y = tensor.tensor3('y')

# Simulate what happens in scan
y2 = y.dimshuffle(1, 0, 2) # (seq_len, 32, 600)
y2 = y2[0, :, :] # (32, 600)
y2 = y2.reshape((32, 600, 1))

z = tensor.batched_dot(x, y2)
f = function([x, y], z)

# (32, 19, 600), (32, 600, 1) -> (32, 19, 1) ?
x_num = np.arange(32 * 19 * 600, dtype=config.floatX).reshape((32, 19, 600))  
y_num = np.arange(7 * 32 * 600, dtype=config.floatX).reshape((32, 7, 600)) 

a = f(x_num, y_num)
```
Just a note, @lamblin did a PR that fix this, the PR just miss test, so you should be able to use it for now.

Resolving patch:
diff --git a/theano/gpuarray/opt.py b/theano/gpuarray/opt.py
--- a/theano/gpuarray/opt.py
+++ b/theano/gpuarray/opt.py
@@ -1208,9 +1208,31 @@ def local_gpua_gemmbatch(op, context_name, inputs, outputs):
     if inputs[0].dtype not in ['float32', 'float64']:
         return
     a, b = inputs
-    c = tensor.AllocEmpty(a.dtype)(a.shape[0], a.shape[1], b.shape[2])
-    return gpugemmbatch_no_inplace(c, np.asarray(1.0, dtype=a.dtype),
-                                   a, b, np.asarray(0.0, dtype=a.dtype))
+    # Since GpuGemmBatch only supports 3D inputs and output,
+    # we need to add broadcastable dims to the inputs, and drop
+    # them from outputs
+    output_dims = [0, 1, 2]
+    if a.ndim == 2:
+        a = GpuDimShuffle(a.broadcastable, (0, 1, 'x'))(a)
+        del output_dims[1]
+    if b.ndim == 2:
+        b = GpuDimShuffle(b.broadcastable, (0, 'x', 1))(b)
+        del output_dims[-1]
+    # In case of mismatched dtypes, we also have to upcast
+    out_dtype = outputs[0].dtype
+    if a.dtype != out_dtype or b.dtype != out_dtype:
+        gpu_cast_op = GpuElemwise(Cast(Scalar(out_dtype)))
+        if a.dtype != out_dtype:
+            a = gpu_cast_op(a)
+        if b.dtype != out_dtype:
+            b = gpu_cast_op(b)
+
+    c = tensor.AllocEmpty(out_dtype)(a.shape[0], a.shape[1], b.shape[2])
+    out = gpugemmbatch_no_inplace(c, np.asarray(1.0, dtype=out_dtype),
+                                  a, b, np.asarray(0.0, dtype=out_dtype))
+    if len(output_dims) != 3:
+        out = GpuDimShuffle(out.broadcastable, output_dims)(out)
+    return out
 
 
 @register_opt()


Test patch:
diff --git a/theano/gpuarray/tests/test_blas.py b/theano/gpuarray/tests/test_blas.py
--- a/theano/gpuarray/tests/test_blas.py
+++ b/theano/gpuarray/tests/test_blas.py
@@ -5,6 +5,7 @@
 import numpy as np
 
 import theano
+from theano import config
 from theano import tensor
 from theano.tests import unittest_tools as utt
 from theano.tensor.blas import gemv_inplace, gemm_inplace, _dot22, batched_dot
@@ -13,7 +14,6 @@
 from .. import gpuarray_shared_constructor
 from .config import mode_with_gpu, test_ctx_name
 from .test_basic_ops import makeTester, rand
-
 from ..blas import (gpugemv_inplace, gpugemv_no_inplace,
                     gpugemm_inplace, gpugemm_no_inplace,
                     gpugemmbatch_no_inplace,
@@ -135,6 +135,18 @@ def shared(val):
     )
 
 
+class TestGpuGemmBatchStrided(TestCase):
+    def test0(self):
+        # Reported in https://github.com/Theano/Theano/issues/5730
+        x = tensor.tensor3()
+        y = tensor.tensor3()
+        z = tensor.batched_dot(x, y[:, 0, :, np.newaxis])
+        f = theano.function([x, y], z, mode=mode_with_gpu)
+        x_num = np.arange(32 * 19 * 600, dtype=config.floatX).reshape((32, 19, 600))
+        y_num = np.arange(7 * 32 * 600, dtype=config.floatX).reshape((32, 7, 600))
+        f(x_num, y_num)
+
+
 class TestGpuSger(TestGer):
     def setUp(self):
         self.mode = mode_with_gpu
diff --git a/theano/gpuarray/tests/test_opt.py b/theano/gpuarray/tests/test_opt.py
--- a/theano/gpuarray/tests/test_opt.py
+++ b/theano/gpuarray/tests/test_opt.py
@@ -656,3 +656,27 @@ def test_local_gpua_advanced_incsubtensor():
     w = tensor.set_subtensor(w[tensor.eq(y, 1.0).nonzero()], 100)
     w = tensor.set_subtensor(w[tensor.eq(y, -1.0).nonzero()], 0)
     theano.function([target], w)
+
+
+def test_batched_dot_lifter():
+    # The CPU Op accepts 2D and 3D inputs, as well as mixed dtypes.
+    # Make sure the lifter adds the appropriate dimshuffles and casts
+    rng = np.random.RandomState(utt.fetch_seed())
+
+    def randX(*args):
+        return rng.rand(*args).astype(theano.config.floatX)
+
+    cases = [
+        (randX(3, 5, 7), randX(3, 7)),
+        (randX(3, 5), randX(3, 5, 7)),
+        (randX(3, 5), randX(3, 5)),
+        (rng.rand(3, 5, 7).astype('float32'), randX(3, 7, 9)),
+        (rng.rand(3, 5, 7).astype('float64'), randX(3, 7, 9))]
+    for x_val, y_val in cases:
+        x = tensor.TensorType(broadcastable=[s == 1 for s in x_val.shape],
+                              dtype=x_val.dtype)('x')
+        y = tensor.TensorType(broadcastable=[s == 1 for s in y_val.shape],
+                              dtype=y_val.dtype)('y')
+        z = tensor.batched_dot(x, y)
+        f = theano.function([x, y], z, mode=mode_with_gpu)
+        f(x_val, y_val)