Handling ndim == 2 in TF batch_dot (keras-team#5280)

* Handling ndim == 2 in TF batch_dot

* Adding support for swapped axes when ndim==2

Author: moustaki

keras/backend/tensorflow_backend.py

@@ -831,7 +831,6 @@ def batch_dot(x, y, axes=None):
 
     # Arguments
         x, y: Keras tensors or variables with `ndim >= 2`
-            (With TensorFlow backend, `batch_dot()` only supports `ndim >= 3`)
         axes: list of (or single) int with target dimensions.
             The lengths of `axes[0]` and `axes[1]` should be the same.
 
@@ -870,24 +869,28 @@ def batch_dot(x, y, axes=None):
         (32, 1, 30)
     ```
     """
-    if ndim(x) < 3 or ndim(y) < 3:
-        raise ValueError('Invalid dimensions for batch_dot: ', ndim(x), ndim(y))
     if isinstance(axes, int):
         axes = (axes, axes)
-    if axes is not None:
-        adj_x = None if axes[0] == ndim(x) - 1 else True
-        adj_y = True if axes[1] == ndim(y) - 1 else None
-    else:
-        adj_x = None
-        adj_y = None
-    # TODO: remove later.
-    if hasattr(tf, 'batch_matmul'):
-        try:
-            out = tf.batch_matmul(x, y, adj_a=adj_x, adj_b=adj_y)
-        except TypeError:
-            out = tf.batch_matmul(x, y, adj_x=adj_x, adj_y=adj_y)
+    if ndim(x) == 2 and ndim(y) == 2:
+        if axes[0] == axes[1]:
+            out = tf.reduce_sum(tf.mul(x, y), axes[0])
+        else:
+            out = tf.reduce_sum(tf.mul(tf.transpose(x, [1, 0]), y), axes[1])
     else:
-        out = tf.matmul(x, y, adjoint_a=adj_x, adjoint_b=adj_y)
+        if axes is not None:
+            adj_x = None if axes[0] == ndim(x) - 1 else True
+            adj_y = True if axes[1] == ndim(y) - 1 else None
+        else:
+            adj_x = None
+            adj_y = None
+        # TODO: remove later.
+        if hasattr(tf, 'batch_matmul'):
+            try:
+                out = tf.batch_matmul(x, y, adj_a=adj_x, adj_b=adj_y)
+            except TypeError:
+                out = tf.batch_matmul(x, y, adj_x=adj_x, adj_y=adj_y)
+        else:
+            out = tf.matmul(x, y, adjoint_a=adj_x, adjoint_b=adj_y)
     if ndim(out) == 1:
         out = expand_dims(out, 1)
     return out

tests/keras/backend/test_backends.py

@@ -77,15 +77,26 @@ def test_linear_operations(self):
 
         check_two_tensor_operation('batch_dot', (4, 2, 3), (4, 5, 3),
                                    axes=(2, 2))
+        check_two_tensor_operation('batch_dot', (32, 20), (32, 20), axes=1)
+        check_two_tensor_operation('batch_dot', (32, 20), (32, 20), axes=(1, 1))
         check_single_tensor_operation('transpose', (4, 2))
         check_single_tensor_operation('reverse', (4, 3, 2), axes=1)
         check_single_tensor_operation('reverse', (4, 3, 2), axes=(1, 2))
 
     def test_batch_dot_shape(self):
-        with pytest.raises(ValueError):
-            x_batch = KTF.ones(shape=(32, 20))
-            y_batch = KTF.ones(shape=(32, 20))
-            xy_batch_dot = KTF.batch_dot(x_batch, y_batch, axes=1)
+        x_batch = KTF.ones(shape=(32, 20))
+        y_batch = KTF.ones(shape=(32, 20))
+        xy_batch_dot = KTF.batch_dot(x_batch, y_batch, axes=1)
+        assert_allclose(KTF.eval(xy_batch_dot), np.ones((32, 1)) * 20, atol=1e-05)
+        xy_batch_dot = KTF.batch_dot(x_batch, y_batch, axes=0)
+        assert_allclose(KTF.eval(xy_batch_dot), np.ones((20, 1)) * 32, atol=1e-05)
+        # making sure swapping axes when ndim == 2 works
+        x_batch = KTF.ones(shape=(32, 20))
+        y_batch = KTF.ones(shape=(20, 32))
+        xy_batch_dot = KTF.batch_dot(x_batch, y_batch, axes=(0, 1))
+        assert_allclose(KTF.eval(xy_batch_dot), np.ones((20, 1)) * 32, atol=1e-05)
+        xy_batch_dot = KTF.batch_dot(x_batch, y_batch, axes=(1, 0))
+        assert_allclose(KTF.eval(xy_batch_dot), np.ones((32, 1)) * 20, atol=1e-05)
 
     def test_shape_operations(self):
         # concatenate

tests/keras/test_sequential_model.py

@@ -7,7 +7,7 @@
 
 from keras import backend as K
 from keras.models import Sequential
-from keras.layers.core import Dense, Activation, Merge, Lambda, Reshape
+from keras.layers.core import Dense, Activation, Merge, Lambda
 from keras.utils import np_utils
 from keras.utils.test_utils import get_test_data, keras_test
 from keras.models import model_from_json, model_from_yaml
@@ -287,35 +287,29 @@ def test_merge_dot():
 
     left = Sequential()
     left.add(Dense(input_dim=input_dim, output_dim=nb_hidden))
-    left.add(Reshape((nb_hidden, 1)))
     left.add(Activation('relu'))
 
     right = Sequential()
     right.add(Dense(input_dim=input_dim, output_dim=nb_hidden))
-    right.add(Reshape((nb_hidden, 1)))
     right.add(Activation('relu'))
 
     model = Sequential()
     model.add(Merge([left, right], mode='dot', dot_axes=1))
-    model.add(Reshape((1,)))
     model.add(Dense(nb_class))
     model.add(Activation('softmax'))
 
     model.compile(loss='categorical_crossentropy', optimizer='rmsprop')
 
     left = Sequential()
     left.add(Dense(input_dim=input_dim, output_dim=nb_hidden))
-    left.add(Reshape((nb_hidden, 1)))
     left.add(Activation('relu'))
 
     right = Sequential()
     right.add(Dense(input_dim=input_dim, output_dim=nb_hidden))
-    right.add(Reshape((nb_hidden, 1)))
     right.add(Activation('relu'))
 
     model = Sequential()
     model.add(Merge([left, right], mode='dot', dot_axes=[1, 1]))
-    model.add(Reshape((1,)))
     model.add(Dense(nb_class))
     model.add(Activation('softmax'))