Repro Colab: https://colab.research.google.com/drive/14n21B0Y6WE3EUj28igcAqKT4T_44LULs?usp=sharing
If a parameter in a custom layer is called *_mask, then it will not work with the layer's compute_output_shapes.
The solution is either to rename the parameter or to remove it altogether from the params of compute_output_shapes function (which is allowed even though params of compute_output_shapes are supposed to shadow those of the layer).
This is a very surprising behavior and kind of odd not to be able to name layer arguments in a certain way. If the behavior was introduced for internal reasons (maybe handling of keras masking), see if it still needed and maybe if it could be relaxed to just "mask" instead of *_mask