Repository: keras-team/keras

Issue:
Operators and some `Merge` (e.g. `Add`, `Maximum`) layers do not propagate masks
Masks are not propagated by add/mul operators, nor using some `Merge` layers (e.g. `Add`, `Maximum`). The below also demonstrates that this bug is also partially present in `tf.keras`  (operators fail, but `Add` / `Maximum` work as expected). Update: `Add` and `Maximum` no longer work in tf.keras.

Below is a simple example of learning the parameters of a straight line, `y = m * x + c`, where `(x, y)` are batched into shape `(batch_size, sequence_length, 1)` to mimic common usage e.g. in NLP. Masked examples have `x = y = 0`. We construct a model `y = dense(map_fn(x))`, where we use various `map_fn`s that should not change the capacity of the model (and in some cases should not even affect gradients). A summary of those that work "as expected" (i.e. learn roughly the correct weights) is shown below.

| map_fn | works with `keras` | works with `tf.keras` |
|----------|--------------------------|-----------------------|
| `lamda x: x` | &check;                    | &check;                  |
| `lambda x: x * 1` | &cross;            | &cross;                    |
| `lambda x: x + x - x` | &cross;            | &cross;               |
| `lambda x: x + x` | &cross;            | &cross;                    |
| `lambda x: keras.layers.Add()((x, x)` | &cross;            | &cross;                    |
| `lambda x: keras.layers.Maximum()((x, x)` | &cross;            | &cross;                    |
| `lambda x: AddSelf()(x)` | &check;            | &check;                    |

## example.py

```python
import sys
import typing as tp
import tensorflow as tf
USE_CORE = {"core": True, "tf": False}[sys.argv[1]]

if USE_CORE:
    import keras
    def get_value(x):
        return x.value
else:
    keras = tf.keras
    def get_value(x):
        return x

def get_dataset():
    """Create some masked dummy data, `y = gradient * x + offset`."""
    max_length = 16
    gradient = 5.
    offset = 2.
    batch_size = 64
    examples_per_epoch = 1000

    data_size = batch_size * examples_per_epoch
    x = tf.random.uniform((data_size, max_length))
    mask = (tf.ragged.range(
        tf.random.uniform((data_size,), 1, max_length, dtype=tf.int64)
    ) + 1).to_tensor(shape=(data_size, max_length)) > 0
    x = tf.where(mask, x, tf.zeros_like(x))
    y = tf.where(mask, x * gradient + offset, tf.zeros_like(x))

    x = tf.expand_dims(x, axis=-1)
    y = tf.expand_dims(y, axis=-1)

    dataset = tf.data.Dataset.from_tensor_slices((x, y))

    #### Adding sample_weight like below fixes the issue
    # sample_weight = tf.cast(mask, tf.float32)
    # dataset = tf.data.Dataset.from_tensor_slices((x, y, sample_weight))

    return dataset.batch(batch_size)


def get_model(input_shape, map_fn: tp.Callable):
    """Create model corresponding to `dense(map_fn(x))`."""
    inp = keras.Input(shape=input_shape[1:], batch_size=input_shape[0])
    x = keras.layers.Masking()(inp)
    x = map_fn(x)
    x = keras.layers.Dense(1)(x)
    return keras.Model(inp, x)
    


def build_and_eval(dataset: tf.data.Dataset, map_fn: tp.Callable, desc: str):
    """Build and evaluate a model dense(map_fn(x)) on dataset."""
    print("----------------")
    print(desc)
    print("----------------")
    model = get_model(dataset.element_spec[0].shape, map_fn)
    if hasattr(model.output, "_keras_mask"):
        print("Output has '_keras_mask' attr")
    else:
        print("Output does NOT have a 'keras_mask' attr")
    model.compile(loss="mse", optimizer="sgd")
    model.fit(dataset, epochs=2)
    print([(w.name, get_value(w).numpy().item()) for w in model.weights])


class AddSelf(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.supports_masking = True
    
    def call(self, x):
        return x + x


def main():
    tf.random.set_seed(42)
    dataset = get_dataset()
    
    build_and_eval(dataset, lambda x: x, "x")
    build_and_eval(dataset, lambda x: x * 1, "x * 1")
    build_and_eval(dataset, lambda x: x + x - x, "x + x - x")
    build_and_eval(dataset, lambda x: x + x, "x + x")
    build_and_eval(
        dataset,
        lambda x: keras.layers.Add()((x, x)),
        "keras.layers.Add()((x, x))",
    )
    build_and_eval(
        dataset,
        lambda x: keras.layers.Maximum()((x, x)),
        "keras.layers.Maximum()((x, x))",
    )
    build_and_eval(dataset, lambda x: AddSelf()(x), "AddSelf()(x)")


if __name__ == '__main__':
    main()
```

## Using `keras`

```bash
python example.py core
```

```txt
----------------
x
----------------
Output has '_keras_mask' attr
Epoch 1/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 837us/step - loss: 0.6007
Epoch 2/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 1s 834us/step - loss: 0.0459
[('kernel', 3.0166471004486084), ('bias', 2.8878185749053955)]
----------------
x * 1
----------------
Output does NOT have a 'keras_mask' attr
Epoch 1/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 1s 893us/step - loss: 2.7253
Epoch 2/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 1s 858us/step - loss: 0.4369
[('kernel', 7.281540870666504), ('bias', 0.4343944489955902)]
----------------
x + x - x
----------------
Output does NOT have a 'keras_mask' attr
Epoch 1/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 1s 798us/step - loss: 2.4995
Epoch 2/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 1s 1ms/step - loss: 0.4331
[('kernel', 7.2877912521362305), ('bias', 0.4326573312282562)]
----------------
x + x
----------------
Output does NOT have a 'keras_mask' attr
Epoch 1/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 1ms/step - loss: 1.9006
Epoch 2/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 1s 1ms/step - loss: 0.3997
[('kernel', 3.7006561756134033), ('bias', 0.4010191261768341)]
----------------
keras.layers.Add()((x, x))
----------------
Output does NOT have a 'keras_mask' attr
Epoch 1/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 1ms/step - loss: 1.8395
Epoch 2/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 1s 1ms/step - loss: 0.3997
[('kernel', 3.7006564140319824), ('bias', 0.40101900696754456)]
----------------
keras.layers.Maximum()((x, x))
----------------
Output does NOT have a 'keras_mask' attr
Epoch 1/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 1ms/step - loss: 2.6779
Epoch 2/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 1s 1ms/step - loss: 0.4361
[('kernel', 7.282827377319336), ('bias', 0.4340367913246155)]
----------------
AddSelf()(x)
----------------
Output has '_keras_mask' attr
Epoch 1/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 0.4728
Epoch 2/2
1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 0.0022
[('kernel', 2.5691280364990234), ('bias', 1.8994231224060059)]
```

## Using `tf.keras`

```bash
python example.py tf
```

```txt
----------------
x
----------------
Output has '_keras_mask' attr
Epoch 1/2
1000/1000 [==============================] - 4s 3ms/step - loss: 0.4196
Epoch 2/2
1000/1000 [==============================] - 3s 3ms/step - loss: 0.0084
[('dense/kernel:0', 4.857532501220703), ('dense/bias:0', 2.0762569904327393)]
----------------
x * 1
----------------
Output does NOT have a 'keras_mask' attr
Epoch 1/2
1000/1000 [==============================] - 5s 5ms/step - loss: 1.6487
Epoch 2/2
1000/1000 [==============================] - 4s 4ms/step - loss: 0.4224
[('dense_1/kernel:0', 7.2698187828063965), ('dense_1/bias:0', 0.4376528561115265)]
----------------
x + x - x
----------------
Output does NOT have a 'keras_mask' attr
Epoch 1/2
1000/1000 [==============================] - 6s 6ms/step - loss: 1.3704
Epoch 2/2
1000/1000 [==============================] - 7s 7ms/step - loss: 0.4173
[('dense_2/kernel:0', 7.285768032073975), ('dense_2/bias:0', 0.43321940302848816)]
----------------
x + x
----------------
Output does NOT have a 'keras_mask' attr
Epoch 1/2
1000/1000 [==============================] - 7s 7ms/step - loss: 0.9798
Epoch 2/2
1000/1000 [==============================] - 4s 4ms/step - loss: 0.3998
[('dense_3/kernel:0', 3.7006540298461914), ('dense_3/bias:0', 0.4010206162929535)]
----------------
keras.layers.Add()((x, x))
----------------
Output does NOT have '_keras_mask' attr
Epoch 1/2
1000/1000 [==============================] - 5s 5ms/step - loss: 0.4362
Epoch 2/2
1000/1000 [==============================] - 6s 6ms/step - loss: 3.6697e-05
[('dense_4/kernel:0', 2.498835563659668), ('dense_4/bias:0', 2.0013747215270996)]
----------------
keras.layers.Maximum()((x, x))
----------------
Output doe NOT have a '_keras_mask' attr
Epoch 1/2
1000/1000 [==============================] - 6s 6ms/step - loss: 0.4207
Epoch 2/2
1000/1000 [==============================] - 7s 7ms/step - loss: 0.0084
[('dense_5/kernel:0', 4.857173919677734), ('dense_5/bias:0', 2.076448917388916)]
----------------
AddSelf()(x)
----------------
Output has '_keras_mask' attr
Epoch 1/2
1000/1000 [==============================] - 6s 6ms/step - loss: 0.1189
Epoch 2/2
1000/1000 [==============================] - 6s 6ms/step - loss: 5.4465e-05
[('dense_6/kernel:0', 2.501418113708496), ('dense_6/bias:0', 1.9983253479003906)]
```




Discussion:
Thanks for the report and analysis. We need to fix this for the `Add` and `Maximum` layers (and maybe other merging layers if affected).

At the same time, it is expected that using low-level operators like `-` and `+` will not preserve masks, because masks are a layer-level concept. A mask can be understood or propagated by a `Layer`.

We *could* lower the concept of masking to the `Operation` class, but it would be very challenging to ensure that *all* our operations are compatible with masks. It seems safer to keep the mental model "masks must be processed by layer instances".

> it is expected that using low-level operators like - and + will not preserve masks

Expected by whom? As a user who's fairly familiar with the keras API (and very rarely a contributor up until recently) I consider it extremely surprising that `x + y` gives fundamentally different behaviour to `Add()((x, y))`.

> masks are a layer-level concept

Conceptually, I've always considered masks a property of a tensor. I think everything would be a lot simpler if we had a `MaskedTensor` composite class, but that's a discussion for somewhere else...

> A mask can be understood or propagated by a Layer

I realize `keras-core` has a separate `Operation` class, and up until now I'd thought of them as just parameterless layers - objects that allow for the creation of a pseudo-computation graph. `tf.keras` one the othe hand explicitly creates `LambdaOp` layers. If these layers don't propagate masks, at the very least I would expect a warning about not propagating as with other layers that don't propagate.

>  it would be very challenging to ensure that all our operations are compatible with masks

I don't think compatibility is necessary, but I think raising a `NotImplementedError` or something similar is much more preferable to silently doing unexpected things. I'd much rather use a code base the requires me to explicitly toss out a mask than one that silently does it for me.
> We need to fix this for the Add and Maximum layers (and maybe other merging layers if affected).

They all inherits the base merge layer, which in fact takes care of the mask. I am surprised why isn't it propagated properly. Need to take a look
@jackd , Is the issue still relevant? if so, you may need to rephrase the question about comparison since tf.keras also uses Keras 3 in backend. 
@sachinprasadhs updated OP. Bugs now uniformly present across keras and tf.keras - some working cases no longer work, i.e. behaviour is now worse than when posted.
@jackd ,

Thanks for the report. There were indeed a few bugs that caused two issues:
- no mask was returned on the outputs
- the mask was not taken into account when performing the merge operation (for instance `Minimum` would take the minimum with zero outside of the mask).

I have a PR open to fix those. I will consider that the merge layers are the scope for this issue.

The other discussion points are much broader.

> I think everything would be a lot simpler if we had a MaskedTensor composite class, but that's a discussion for somewhere else...

The concept of `MaskedTensor`s doesn't exist in all backends, so this is not really an option at this time.

> it is expected that using low-level operators like `-` and `+` will not preserve masks

One of the unintuitive things about Keras supporting multiple backends is that `x + y` can behave differently than `ops.add(x, y)` in subtle ways. That is because `+` is handled natively by the backend (Tensorflow, JAX, Torch) without Keras having any visibility into it. Without native backend support for masks, `x + y` won't propagate masks.

> If these layers don't propagate masks, at the very least I would expect a warning about not propagating as with other layers that don't propagate.

This warning exists: https://github.com/keras-team/keras/blob/master/keras/src/layers/layer.py#L879-L885  It didn't fire in this case because the merge layers claim to support masking and do implement `compute_mask`. Simply, the implementation was wrong. Note that not having a mask on the output does not mean that the layer unintentionally lost the mask. For instance, it is expected that `Add()([x, y])` will have no mask if x has a mask and y has no mask.
@hertschuh 

> The concept of MaskedTensors doesn't exist in all backends

So keras only ever plans to implement composite tensors that are implemented in all backends? Getting a bit off topic for this issue, but that seems like an unnecessary restriction to me...

> That is because + is handled natively by the backend (Tensorflow, JAX, Torch)

Isn't [this](https://github.com/keras-team/keras/blob/master/keras/src/backend/common/keras_tensor.py#L130) how add is handled? Looks like any PR that fixes `Add` will also fix `__add__`.

> For instance, it is expected that Add()([x, y]) will have no mask if x has a mask and y has no mask.

That's not at all what I would expect. What is a masked tensor supposed to represent? My understanding is that it's essentially a different representation of tensorflow's `RaggedTensor`. Adding a `RaggedTensor` with shape `[N, ?]` to a regular tensor of shape `[N, M]` results in an error AFAIK. Are there any situations where it makes sense to add a masked and unmasked tensor?
Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=Yes&entry.243948740=https://github.com/keras-team/keras/issues/18416">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=No&entry.243948740=https://github.com/keras-team/keras/issues/18416">No</a>




Patch:
diff --git a/keras/src/layers/layer.py b/keras/src/layers/layer.py
index e668ca342080..99856bbea04c 100644
--- a/keras/src/layers/layer.py
+++ b/keras/src/layers/layer.py
@@ -870,12 +870,14 @@ def maybe_convert(x):
             # Set masks on outputs,
             # provided only the first positional input arg and its mask.
             # TODO: consider extending this to all args and kwargs.
-            previous_mask = getattr(call_spec.first_arg, "_keras_mask", None)
+            previous_mask = tree.map_structure(
+                lambda x: getattr(x, "_keras_mask", None), call_spec.first_arg
+            )
             if self.supports_masking:
                 self._set_mask_metadata(
                     call_spec.first_arg, outputs, previous_mask
                 )
-            elif previous_mask is not None:
+            elif any(m is not None for m in tree.flatten(previous_mask)):
                 warnings.warn(
                     f"Layer '{self.name}' (of type {self.__class__.__name__}) "
                     "was passed an input with a mask attached to it. "
diff --git a/keras/src/layers/merging/base_merge.py b/keras/src/layers/merging/base_merge.py
index 69591e8cd074..3a56d76e4ce8 100644
--- a/keras/src/layers/merging/base_merge.py
+++ b/keras/src/layers/merging/base_merge.py
@@ -20,6 +20,43 @@ def __init__(self, **kwargs):
     def _merge_function(self, inputs):
         raise NotImplementedError
 
+    def _apply_merge_op_and_or_mask(self, op_fn, inputs):
+        """Merge a set of inputs by applying `op_fn` and ORing the masks.
+
+        We use this for `Minimum` and `Maximum` as it handles the fact that
+        there is no identity element. If applicable, the mask obtained by ORing
+        all masks is set on the output.
+
+        Args:
+            op_fn: binary operation to apply to tensor pair.
+            inputs: array of tensors to apply operation on.
+        """
+        output = None
+        output_mask = None
+
+        for x in inputs:
+            mask = getattr(x, "_keras_mask", None)
+            if mask is not None:
+                mask = ops.broadcast_to(ops.expand_dims(mask, -1), ops.shape(x))
+            if output is None:
+                output = x
+                output_mask = mask
+                continue
+            if mask is not None:
+                x = ops.where(mask, x, output)
+            if output_mask is not None:
+                output = ops.where(output_mask, output, x)
+            if mask is not None and output_mask is not None:
+                output_mask = ops.logical_or(output_mask, mask)
+            else:
+                output_mask = None
+            output = op_fn(output, x)
+
+        if output_mask is not None:
+            output_mask = ops.any(output_mask, axis=-1, keepdims=False)
+            output._keras_mask = output_mask
+        return output
+
     def _compute_elemwise_op_output_shape(self, shape1, shape2):
         """Computes the shape of the resultant of an elementwise operation.
 
@@ -231,10 +268,14 @@ def compute_mask(self, inputs, mask=None):
                 f"Received: inputs={inputs} of length {len(inputs)}, and "
                 f"mask={mask} of length {len(mask)}"
             )
-        if all(m is None for m in mask):
+        # Default implementation does an OR between the masks, which works
+        # for `Add`, `Subtract`, `Average`, `Maximum`, `Minimum`, `Multiply`.
+        if any(m is None for m in mask):
             return None
-        masks = [ops.expand_dims(m, axis=0) for m in mask if m is not None]
-        return ops.all(ops.concatenate(masks, axis=0), axis=0, keepdims=False)
+        output_mask = mask[0]
+        for m in mask[1:]:
+            output_mask = ops.logical_or(output_mask, m)
+        return output_mask
 
     def get_config(self):
         return super().get_config()
diff --git a/keras/src/layers/merging/concatenate.py b/keras/src/layers/merging/concatenate.py
index 9c1d26ae8d5f..7e240786ac3e 100644
--- a/keras/src/layers/merging/concatenate.py
+++ b/keras/src/layers/merging/concatenate.py
@@ -145,11 +145,15 @@ def compute_mask(self, inputs, mask=None):
                 masks.append(ops.ones_like(input_i, dtype="bool"))
             elif mask_i.ndim < input_i.ndim:
                 # Mask is smaller than the input, expand it
-                masks.append(ops.expand_dims(mask_i, axis=-1))
+                masks.append(
+                    ops.broadcast_to(
+                        ops.expand_dims(mask_i, axis=-1), ops.shape(input_i)
+                    )
+                )
             else:
                 masks.append(mask_i)
         concatenated = ops.concatenate(masks, axis=self.axis)
-        return ops.all(concatenated, axis=-1, keepdims=False)
+        return ops.any(concatenated, axis=-1, keepdims=False)
 
     def get_config(self):
         config = {"axis": self.axis}
diff --git a/keras/src/layers/merging/maximum.py b/keras/src/layers/merging/maximum.py
index 47734a8470d5..3072ecb625a9 100644
--- a/keras/src/layers/merging/maximum.py
+++ b/keras/src/layers/merging/maximum.py
@@ -31,10 +31,7 @@ class Maximum(Merge):
     """
 
     def _merge_function(self, inputs):
-        output = inputs[0]
-        for i in range(1, len(inputs)):
-            output = ops.maximum(output, inputs[i])
-        return output
+        return self._apply_merge_op_and_or_mask(ops.maximum, inputs)
 
 
 @keras_export("keras.layers.maximum")
diff --git a/keras/src/layers/merging/minimum.py b/keras/src/layers/merging/minimum.py
index 19137f08a5b3..dad5997ef656 100644
--- a/keras/src/layers/merging/minimum.py
+++ b/keras/src/layers/merging/minimum.py
@@ -31,10 +31,7 @@ class Minimum(Merge):
     """
 
     def _merge_function(self, inputs):
-        output = inputs[0]
-        for i in range(1, len(inputs)):
-            output = ops.minimum(output, inputs[i])
-        return output
+        return self._apply_merge_op_and_or_mask(ops.minimum, inputs)
 
 
 @keras_export("keras.layers.minimum")
diff --git a/keras/src/layers/merging/multiply.py b/keras/src/layers/merging/multiply.py
index d908429d1c5b..72ed99bf038a 100644
--- a/keras/src/layers/merging/multiply.py
+++ b/keras/src/layers/merging/multiply.py
@@ -31,9 +31,29 @@ class Multiply(Merge):
     """
 
     def _merge_function(self, inputs):
-        output = inputs[0]
-        for i in range(1, len(inputs)):
-            output = ops.multiply(output, inputs[i])
+        masks = [getattr(x, "_keras_mask", None) for x in inputs]
+        has_output_mask = all(mask is not None for mask in masks)
+        output = None
+        output_mask = None
+
+        for x, mask in zip(inputs, masks):
+            if mask is not None:
+                mask = ops.broadcast_to(ops.expand_dims(mask, -1), ops.shape(x))
+                # Replace 0s with 1s outside of mask.
+                x = ops.where(mask, x, ops.cast(1, x.dtype))
+                if has_output_mask:
+                    output_mask = (
+                        mask
+                        if output_mask is None
+                        else ops.logical_or(output_mask, mask)
+                    )
+            output = x if output is None else ops.multiply(output, x)
+
+        if has_output_mask:
+            # Replace 1s with 0s outside of mask per standard masking rules.
+            output = ops.where(output_mask, output, ops.cast(0, output.dtype))
+            output_mask = ops.any(output_mask, axis=-1, keepdims=False)
+            output._keras_mask = output_mask
         return output
 
 


Test patch:
diff --git a/keras/src/layers/merging/merging_test.py b/keras/src/layers/merging/merging_test.py
index 1419dc855b83..b9df8264e225 100644
--- a/keras/src/layers/merging/merging_test.py
+++ b/keras/src/layers/merging/merging_test.py
@@ -125,14 +125,14 @@ def test_correctness_static(
         self.assertEqual(res.shape, expected_output_shape)
         self.assertAllClose(res, x3, atol=1e-4)
         self.assertIsNone(layer.compute_mask([input_1, input_2], [None, None]))
+        self.assertIsNone(layer.compute_mask([x1, x2], [None, None]))
         if not skip_mask_test:
+            mask1 = np.ones(input_shape[:-1], dtype=np.bool_)
+            mask2 = np.ones(input_shape[:-1], dtype=np.bool_)
             self.assertTrue(
                 np.all(
                     backend.convert_to_numpy(
-                        layer.compute_mask(
-                            [input_1, input_2],
-                            [backend.Variable(x1), backend.Variable(x2)],
-                        )
+                        layer.compute_mask([x1, x2], [mask1, mask2])
                     )
                 )
             )
@@ -234,6 +234,111 @@ def test_dot_higher_dim(self):
         c = layers.Dot(axes=(-2, -1))([a, b])
         self.assertEqual(backend.standardize_shape(c.shape), (1, 2, 1, 2))
 
+    def test_add_with_mask(self):
+        mask = layers.Masking()
+        x1 = mask(backend.convert_to_tensor([[[0, 0], [1, 2], [0, 0], [3, 4]]]))
+        x2 = backend.convert_to_tensor([[[0, 0], [0, 0], [1, 2], [3, 4]]])
+
+        output = layers.Add()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [1, 2], [1, 2], [6, 8]]])
+        self.assertIsNone(getattr(output, "_keras_mask", None))
+
+        x2 = mask(x2)
+        output = layers.Add()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [1, 2], [1, 2], [6, 8]]])
+        self.assertAllClose(output._keras_mask, [[0, 1, 1, 1]])
+
+    def test_subtract_with_mask(self):
+        mask = layers.Masking()
+        x1 = mask(backend.convert_to_tensor([[[0, 0], [1, 2], [0, 0], [3, 4]]]))
+        x2 = backend.convert_to_tensor([[[0, 0], [0, 0], [1, 2], [3, 4]]])
+
+        output = layers.Subtract()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [1, 2], [-1, -2], [0, 0]]])
+        self.assertIsNone(getattr(output, "_keras_mask", None))
+
+        x2 = mask(x2)
+        output = layers.Subtract()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [1, 2], [-1, -2], [0, 0]]])
+        self.assertAllClose(output._keras_mask, [[0, 1, 1, 1]])
+
+    def test_average_with_mask(self):
+        mask = layers.Masking()
+        x1 = mask(backend.convert_to_tensor([[[0, 0], [1, 2], [0, 0], [3, 4]]]))
+        x2 = backend.convert_to_tensor([[[0, 0], [0, 0], [1, 2], [3, 4]]])
+
+        output = layers.Average()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [0.5, 1], [0.5, 1], [3, 4]]])
+        self.assertIsNone(getattr(output, "_keras_mask", None))
+
+        x2 = mask(x2)
+        output = layers.Average()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [0.5, 1], [0.5, 1], [3, 4]]])
+        self.assertAllClose(output._keras_mask, [[0, 1, 1, 1]])
+
+    def test_multiply_with_mask(self):
+        mask = layers.Masking()
+        x1 = mask(backend.convert_to_tensor([[[0, 0], [1, 2], [0, 0], [3, 4]]]))
+        x2 = backend.convert_to_tensor([[[0, 0], [0, 0], [1, 2], [3, 4]]])
+
+        output = layers.Multiply()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [0, 0], [1, 2], [9, 16]]])
+        self.assertIsNone(getattr(output, "_keras_mask", None))
+
+        x2 = mask(x2)
+        output = layers.Multiply()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [1, 2], [1, 2], [9, 16]]])
+        self.assertAllClose(output._keras_mask, [[0, 1, 1, 1]])
+
+    def test_maximum_with_mask(self):
+        mask = layers.Masking()
+        x1 = mask(
+            backend.convert_to_tensor([[[0, 0], [-1, -2], [0, 0], [-3, -4]]])
+        )
+        x2 = backend.convert_to_tensor([[[0, 0], [0, 0], [-1, -2], [-3, -4]]])
+
+        output = layers.Maximum()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [0, 0], [-1, -2], [-3, -4]]])
+        self.assertIsNone(getattr(output, "_keras_mask", None))
+
+        x2 = mask(x2)
+        output = layers.Maximum()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [-1, -2], [-1, -2], [-3, -4]]])
+        self.assertAllClose(output._keras_mask, [[0, 1, 1, 1]])
+
+    def test_minimum_with_mask(self):
+        mask = layers.Masking()
+        x1 = mask(backend.convert_to_tensor([[[0, 0], [1, 2], [0, 0], [3, 4]]]))
+        x2 = backend.convert_to_tensor([[[0, 0], [0, 0], [1, 2], [3, 4]]])
+
+        output = layers.Minimum()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [0, 0], [1, 2], [3, 4]]])
+        self.assertIsNone(getattr(output, "_keras_mask", None))
+
+        x2 = mask(x2)
+        output = layers.Minimum()([x1, x2])
+        self.assertAllClose(output, [[[0, 0], [1, 2], [1, 2], [3, 4]]])
+        self.assertAllClose(output._keras_mask, [[0, 1, 1, 1]])
+
+    def test_concatenate_with_mask(self):
+        mask = layers.Masking()
+        x1 = mask(backend.convert_to_tensor([[[0, 0], [1, 2], [0, 0], [3, 4]]]))
+        x2 = backend.convert_to_tensor([[[0, 0], [0, 0], [1, 2], [3, 4]]])
+
+        output = layers.Concatenate(axis=1)([x1, x2])
+        self.assertAllClose(
+            output,
+            [[[0, 0], [1, 2], [0, 0], [3, 4], [0, 0], [0, 0], [1, 2], [3, 4]]],
+        )
+        self.assertAllClose(output._keras_mask, [[0, 1, 0, 1, 1, 1, 1, 1]])
+
+        output = layers.Concatenate(axis=2)([x1, x2])
+        self.assertAllClose(
+            output,
+            [[[0, 0, 0, 0], [1, 2, 0, 0], [0, 0, 1, 2], [3, 4, 3, 4]]],
+        )
+        self.assertAllClose(output._keras_mask, [[1, 1, 1, 1]])
+
     @parameterized.named_parameters(TEST_PARAMETERS)
     @pytest.mark.skipif(
         not backend.SUPPORTS_SPARSE_TENSORS,