Merge layers do not pass masks to the `compute_mask` function

[jbed94]

After Keras & TensorFlow upgrade to 3.3 and 2.16.1 I encountered a bug with mask propagation in merge layers. A base Layer class in __call__ function takes a first argument and treats it as a previous_mask (link to source). It produces an incorrect mask propagation in Merge layers where arguments are passed as a single list of tensors (and list doesn't have _keras_mask attribute).

Example code to reproduce unexpected behavior:

from keras import layers, random, ops, Model

a = random.uniform([1, 2, 256], dtype="float32")
b = random.uniform([1, 2, 256], dtype="float32")

a._keras_mask = ops.convert_to_tensor([[True, False]])
b._keras_mask = ops.convert_to_tensor([[True, False]])

add_op = layers.Add()

x = add_op([a, b])
expected_mask = add_op.compute_mask([a, b], [a._keras_mask, b._keras_mask])

assert hasattr(x, "_keras_mask")

I see, that the code part is marked as TODO, are there any plans to fix it?

=== Edit ===

Similar code works with Keras 2.15 and TensorFlow 2.15:

import tensorflow as tf

a = tf.random.uniform([1, 2, 256], dtype="float32")
b = tf.random.uniform([1, 2, 256], dtype="float32")

a._keras_mask = tf.convert_to_tensor([[True, False]])
b._keras_mask = tf.convert_to_tensor([[True, False]])

add_op = tf.keras.layers.Add()

x = add_op([a, b])
expected_mask = add_op.compute_mask([a, b], [a._keras_mask, b._keras_mask])

assert hasattr(x, "_keras_mask")
assert tf.reduce_all(x._keras_mask == expected_mask)

[Grvzard]

The issue has been fixed in the main branch, but it hasn't been released yet. Related PR: #19685

[sachinprasadhs]

@jbed94 , You can use the Keras nightly till the next release, here is the working Gist using Keras-nightly

[github-actions[bot]]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

[github-actions[bot]]

This issue was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further.

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No