Mask equal None when using TimeDistributed Tensorflow 2 - deep-learning

Question

I have model with encoding layer. I want to use TimeDistributed to encoding layer in my model. But when I run the model with TimeDistributed, it stuck at encoding layer in mask = tf.cast(mask, dtype=self.dtype) and it mean my mask will always None when using TimeDistributed layer. How can I apply TimeDistributed layer to my encoding layer?
The encoding layer code:
class Encoding(Layer):
def __init__(self,
**kwargs):
super().__init__(**kwargs)
def build(self, input_shape):
self.encoding = self.add_weight(shape=(input_shape[-2], input_shape[-1]),
trainable=True,
initializer=tf.initializers.Ones(),
constraint=None,
dtype=self.dtype, name='encoding')
super().build(input_shape)
def call(self, inputs, mask=None):
mask = tf.cast(mask, dtype=self.dtype)
mask = tf.expand_dims(mask, axis=-1)
return tf.reduce_sum(mask * self.encoding * inputs, axis=-2)
def compute_mask(self, inputs, mask=None):
if mask is None:
return None
return tf.reduce_any(mask, axis=-1)
def compute_output_shape(self, input_shape):
return (input_shape[0], input_shape[-1])
if I build model with encoding in it like this:
encoding = Encoding(name='encoded')
story_encoded = encoding(story_embedded)
query_encoded = encoding(query_embedded)
It can run without error. But of I use TimeDistributed to my encoding:
encoding = Encoding(name='encoded')
story_encoded = TimeDistributed(encoding, name='story_encoding')(story_embedded)
query_encoded = TimeDistributed(encoding, name='query_encoding')(query_embedded)
It raise error like this:
ValueError Traceback (most recent call last)
<ipython-input-15-6c1bbfdcab4e> in <module>()
42 shuffle=True,
43 callbacks=callbacks,
---> 44 verbose=1
45 )
46
1 frames
/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
65 except Exception as e: # pylint: disable=broad-except
66 filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67 raise e.with_traceback(filtered_tb) from None
68 finally:
69 del filtered_tb
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py in autograph_handler(*args, **kwargs)
1145 except Exception as e: # pylint:disable=broad-except
1146 if hasattr(e, "ag_error_metadata"):
-> 1147 raise e.ag_error_metadata.to_exception(e)
1148 else:
1149 raise
ValueError: in user code:
File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1021, in train_function *
return step_function(self, iterator)
File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1010, in step_function **
outputs = model.distribute_strategy.run(run_step, args=(data,))
File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1000, in run_step **
outputs = model.train_step(data)
File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 859, in train_step
y_pred = self(x, training=True)
File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler
raise e.with_traceback(filtered_tb) from None
ValueError: Exception encountered when calling layer "encoded" (type Encoding).
in user code:
File "<ipython-input-4-9d7376af8f80>", line 26, in call *
mask = tf.cast(mask, dtype=self.dtype)
ValueError: None values not supported.
Call arguments received:
• inputs=tf.Tensor(shape=(30, 3, 80), dtype=float32)
• mask=None