Support for masking in merged layers
diff --git a/keras/backend/tensorflow_backend.py b/keras/backend/tensorflow_backend.py
index 2616ec6d7b1..833d31019c1 100644
--- a/keras/backend/tensorflow_backend.py
+++ b/keras/backend/tensorflow_backend.py
@@ -339,6 +339,17 @@ def any(x, axis=None, keepdims=False):
     return tf.cast(x, tf.uint8)
 
 
+def all(x, axis=None, keepdims=False):
+    '''Bitwise reduction (logical AND).
+
+    Returns an uint8 tensor
+    '''
+    axis = _normalize_axis(axis, ndim(x))
+    x = tf.cast(x, tf.bool)
+    x = tf.reduce_all(x, reduction_indices=axis, keep_dims=keepdims)
+    return tf.cast(x, tf.uint8)
+
+
 def argmax(x, axis=-1):
     '''Returns the index of the maximum value
     along a tensor axis.
diff --git a/keras/backend/theano_backend.py b/keras/backend/theano_backend.py
index 8c391e4bb1b..13a21a80493 100644
--- a/keras/backend/theano_backend.py
+++ b/keras/backend/theano_backend.py
@@ -202,6 +202,12 @@ def any(x, axis=None, keepdims=False):
     return T.any(x, axis=axis, keepdims=keepdims)
 
 
+def all(x, axis=None, keepdims=False):
+    '''Bitwise reduction (logical AND).
+    '''
+    return T.all(x, axis=axis, keepdims=keepdims)
+
+
 def argmax(x, axis=-1):
     return T.argmax(x, axis=axis, keepdims=False)
 
diff --git a/keras/engine/topology.py b/keras/engine/topology.py
index 30942286a4e..a26e43982f6 100644
--- a/keras/engine/topology.py
+++ b/keras/engine/topology.py
@@ -1073,9 +1073,12 @@ class Merge(Layer):
         tensor_indices: optional list of indices of output tensors
             to consider for merging
             (in case some input layer node returns multiple tensors).
+        output_mask: mask or lambda/function to compute the output mask (only
+            if merge mode is a lambda/function). If the latter case, it should
+            take as input a list of masks and return a single mask.
     '''
     def __init__(self, layers=None, mode='sum', concat_axis=-1,
-                 dot_axes=-1, output_shape=None,
+                 dot_axes=-1, output_shape=None, output_mask=None,
                  node_indices=None, tensor_indices=None, name=None):
         self.layers = layers
         self.mode = mode
@@ -1085,6 +1088,7 @@ def __init__(self, layers=None, mode='sum', concat_axis=-1,
             self.dot_axes = [self.dot_axes, ] * 2
         self._output_shape = output_shape
         self.node_indices = node_indices
+        self._output_mask = output_mask
 
         # layer parameters
         self.inbound_nodes = []
@@ -1093,7 +1097,7 @@ def __init__(self, layers=None, mode='sum', concat_axis=-1,
         self.regularizers = []
         self.trainable_weights = []
         self.non_trainable_weights = []
-        self.supports_masking = False
+        self.supports_masking = True
         self.uses_learning_phase = False
         self.input_spec = None  # compatible with whatever
         if not name:
@@ -1317,13 +1321,30 @@ def get_output_shape_for(self, input_shape):
         elif self.mode == 'cos':
             return (input_shapes[0][0], 1)
 
-    def compute_mask(self, input, mask=None):
-        '''TODO: add mask merging support
-        '''
-        if mask is not None and any(mask):
-            raise Exception('Merge does not support masking, ' +
-                            'but was passed an input mask: ' + str(mask))
-        return None
+    def compute_mask(self, inputs, mask=None):
+        if mask is None or not any([m is not None for m in mask]):
+            return None
+
+        assert hasattr(mask, '__len__') and len(mask) == len(inputs)
+
+        if self.mode in ['sum', 'mul', 'ave']:
+            masks = [K.expand_dims(m, 0) for m in mask if m is not None]
+            return K.all(K.concatenate(masks, axis=0), axis=0, keepdims=False)
+        elif self.mode == 'concat':
+            masks = [K.ones_like(inputs[i][:-1]) if m is None else m for i, m in zip(inputs, mask)]
+            expanded_dims = [K.expand_dims(m) for m in masks]
+            concatenated = K.concatenate(expanded_dims, axis=self.concat_axis)
+            return K.all(concatenated, axis=-1, keepdims=False)
+        elif self.mode in ['cos', 'dot']:
+            return None
+        elif hasattr(self.mode, '__call__'):
+            if hasattr(self._output_mask, '__call__'):
+                return self._output_mask(mask)
+            else:
+                return self._output_mask
+        else:
+            # this should have been caught earlier
+            raise Exception('Invalid merge mode: {}'.format(self.mode))
 
     def get_config(self):
         py3 = sys.version_info[0] == 3
@@ -1388,7 +1409,7 @@ def from_config(cls, config):
 
 
 def merge(inputs, mode='sum', concat_axis=-1,
-          dot_axes=-1, output_shape=None, name=None):
+          dot_axes=-1, output_shape=None, output_mask=None, name=None):
     '''Functional merge, to apply to Keras tensors (NOT layers).
     Returns a Keras tensor.
 
@@ -1437,6 +1458,7 @@ def merge(inputs, mode='sum', concat_axis=-1,
                             concat_axis=concat_axis,
                             dot_axes=dot_axes,
                             output_shape=output_shape,
+                            output_mask=output_mask,
                             node_indices=node_indices,
                             tensor_indices=tensor_indices,
                             name=name)
@@ -1446,6 +1468,7 @@ def merge(inputs, mode='sum', concat_axis=-1,
                             concat_axis=concat_axis,
                             dot_axes=dot_axes,
                             output_shape=output_shape,
+                            output_mask=output_mask,
                             name=name)
         return merge_layer(inputs)
 
diff --git a/keras/layers/core.py b/keras/layers/core.py
index b7c20c4f7bc..75f83f7cd8b 100644
--- a/keras/layers/core.py
+++ b/keras/layers/core.py
@@ -402,6 +402,8 @@ def antirectifier_output_shape(input_shape):
     def __init__(self, function, output_shape=None, arguments={}, **kwargs):
         self.function = function
         self.arguments = arguments
+        self.supports_masking = True
+
         if output_shape is None:
             self._output_shape = None
         elif type(output_shape) in {tuple, list}:
diff --git a/tests/keras/backend/test_backends.py b/tests/keras/backend/test_backends.py
index 465d1dae455..fb4b2647464 100644
--- a/tests/keras/backend/test_backends.py
+++ b/tests/keras/backend/test_backends.py
@@ -139,6 +139,9 @@ def test_elementwise_operations(self):
         # does not work yet, wait for bool <-> int casting in TF (coming soon)
         # check_single_tensor_operation('any', (4, 2))
         # check_single_tensor_operation('any', (4, 2), axis=1, keepdims=True)
+        #
+        # check_single_tensor_operation('any', (4, 2))
+        # check_single_tensor_operation('any', (4, 2), axis=1, keepdims=True)
 
         check_single_tensor_operation('argmax', (4, 2))
         check_single_tensor_operation('argmax', (4, 2), axis=1)
diff --git a/tests/keras/layers/test_core.py b/tests/keras/layers/test_core.py
index 281e5ce61af..bbeab907865 100644
--- a/tests/keras/layers/test_core.py
+++ b/tests/keras/layers/test_core.py
@@ -1,6 +1,5 @@
 import pytest
 import numpy as np
-from numpy.testing import assert_allclose
 
 from keras import backend as K
 from keras.layers import core
@@ -84,6 +83,60 @@ def fn_output_shape(tup):
     model.compile('rmsprop', 'mse')
 
 
+def test_merge_mask_2d():
+    from keras.layers import Input, merge, Masking
+    from keras.models import Model
+
+    rand = lambda *shape: np.asarray(np.random.random(shape) > 0.5, dtype='int32')
+
+    # inputs
+    input_a = Input(shape=(3,))
+    input_b = Input(shape=(3,))
+
+    # masks
+    masked_a = Masking(mask_value=0)(input_a)
+    masked_b = Masking(mask_value=0)(input_b)
+
+    # two different types of merging
+    merged_sum = merge([masked_a, masked_b], mode='sum')
+    merged_concat = merge([masked_a, masked_b], mode='concat', concat_axis=1)
+
+    # test sum
+    model_sum = Model([input_a, input_b], [merged_sum])
+    model_sum.compile(loss='mse', optimizer='sgd')
+    model_sum.fit([rand(2,3), rand(2,3)], [rand(2,3)], nb_epoch=1)
+
+    # test concatenation
+    model_concat = Model([input_a, input_b], [merged_concat])
+    model_concat.compile(loss='mse', optimizer='sgd')
+    model_concat.fit([rand(2,3), rand(2,3)], [rand(2,6)], nb_epoch=1)
+
+
+def test_merge_mask_3d():
+    from keras.layers import Input, merge, Embedding, SimpleRNN
+    from keras.models import Model
+
+    rand = lambda *shape: np.asarray(np.random.random(shape) > 0.5, dtype='int32')
+
+    # embeddings
+    input_a = Input(shape=(3,), dtype='int32')
+    input_b = Input(shape=(3,), dtype='int32')
+    embedding = Embedding(3, 4, mask_zero=True)
+    embedding_a = embedding(input_a)
+    embedding_b = embedding(input_b)
+
+    # rnn
+    rnn = SimpleRNN(3, return_sequences=True)
+    rnn_a = rnn(embedding_a)
+    rnn_b = rnn(embedding_b)
+
+    # concatenation
+    merged_concat = merge([rnn_a, rnn_b], mode='concat', concat_axis=-1)
+    model = Model([input_a, input_b], [merged_concat])
+    model.compile(loss='mse', optimizer='sgd')
+    model.fit([rand(2,3), rand(2,3)], [rand(2,3,6)])
+
+
 def test_dropout():
     layer_test(core.Dropout,
                kwargs={'p': 0.5},

I've seen a couple other people who seem like they need masking in merged layers. As far as I can tell, masking is only required for the `'sum'` and `'ave'` modes (and lambda functions as well probably, but those are a case of their own). I've only tested this fork with Theano + Python 2.7, but it seems to be working.

I did notice some weird behavior with the `Masking` layer. If given this data,

```
input a data:
[[ 1.  1.  1.]
 [ 1.  1.  1.]]

input b data:
[[ 1.  1.  0.]
 [ 0.  0.  0.]]
```

for some model that adds the inputs but masks 0's in the second input,

```
input_a = Input(shape=(3,), dtype='int32')
input_b = Input(shape=(3,), dtype='int32')
masked_b = Masking(mask_value=0)(input_b)
merged = merge([input_a, masked_b], mode='sum')
```

it gives

```
[[ 3.  3.  2.]
 [ 0.  0.  0.]]
```

Is this the correct behavior? It was my best guess at what it should be doing. Also, I haven't tried this with `mask_zeros` in the `Embedding` layer, which seems like the most common use case.

Also, I'm not sure how to go about adding support for lambda functions, since this also seems useful. I thought maybe adding an `output_mask` parameter that takes a list of (normalized) masks and returns the output mask, but this seems clunky.
