Merge pull request #1020 from caglar/batched_dot22

theano/tensor/basic.py

@@ -4370,6 +4370,27 @@ def set_subtensor(x, y, inplace=False,
     return inc_subtensor(x, y, inplace, set_instead_of_inc=True,
             tolerate_inplace_aliasing=tolerate_inplace_aliasing)
 
+def batched_dot(x, y):
+    """
+    :param x: A Tensor with sizes e.g.: for  3D (dim1, dim3, dim2)
+    :param y: A Tensor with sizes e.g.: for 3D (dim1, dim2, dim4)
+    This function computes the dot product between the two tensors, by iterating
+    over the first dimension using scan.
+    Returns a tensor of size e.g. if it is 3D: (dim1, dim3, dim4)
+    Example:
+    >>> first = T.tensor3('first')
+    >>> second = T.tensor3('second')
+    >>> result = batched_dot(first, second)
+    :note:  This is a subset of numpy.einsum, but we do not provide it for now.
+    But numpy einsum is slower then dot or tensordot:
+    http://mail.scipy.org/pipermail/numpy-discussion/2012-October/064259.html
+    """
+    result, updates = theano.scan(fn=lambda x_mat, y_mat:
+            theano.tensor.dot(x_mat, y_mat),
+            outputs_info=None,
+            sequences=[x, y],
+            non_sequences=None)
+    return result
 
 def inc_subtensor(x, y, inplace=False, set_instead_of_inc=False,
         tolerate_inplace_aliasing=False):

theano/tensor/tests/test_basic.py

@@ -1893,6 +1893,27 @@ def _approx_eq(a, b, eps=1.0e-4):
     return  True
 _approx_eq.debug = 0
 
+def test_batched_dot():
+    first = theano.tensor.tensor3("first")
+    second = theano.tensor.tensor3("second")
+    output = theano.tensor.basic.batched_dot(first, second)
+    first_val = numpy.random.rand(10, 10, 20)
+    second_val = numpy.random.rand(10, 20, 5)
+    result_fn = theano.function([first, second], output)
+    result = result_fn(first_val, second_val)
+    assert result.shape[0] == first_val.shape[0]
+    assert result.shape[1] == first_val.shape[1]
+    assert result.shape[2] == second_val.shape[2]
+
+    first_mat = theano.tensor.dmatrix("first")
+    second_mat = theano.tensor.dmatrix("second")
+    output = theano.tensor.basic.batched_dot(first_mat, second_mat)
+    first_mat_val = numpy.random.rand(10, 10)
+    second_mat_val = numpy.random.rand(10, 10)
+    result_fn = theano.function([first_mat, second_mat], output)
+    result = result_fn(first_mat_val, second_mat_val)
+
+    assert result.shape[0] == first_val.shape[0]
 
 def test_tensor_values_eq_approx():
     #test, inf, -inf and nan equal themself