tensorflow: The inheriting keras.layers.Layer does not call a compute_output_shape after switching to tf.keras from keras

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow):
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary): pip
• TensorFlow version (use command below): 1.14, 1.15 and 2.0 (gpu)
• Keras version: 2.2.4 and 2.2.4-tf
• Python version: 3.6.1
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: 10/6.7.4
• GPU model and memory: RTX 2060 6GB

You can also obtain the TensorFlow version with: 1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)" unknown 1.14.0

After switching to tf.keras from keras my project witch use Mask_RCNN does not work. I noticed that TimeDistributed returns a different tensor shape in keras (1.) and tf.keras (2.): 1. Tensor("mrcnn_class_bn1/Reshape_1:0", shape=(?, 1000, 1, 1, 256), dtype=float32) 2. Tensor("mrcnn_class_bn1/Reshape_1:0", shape=(?, ?, 1, 1, 256), dtype=float32)

x = PyramidROIAlign([pool_size, pool_size], feature_levels,
                        name="roi_align_classifier")([rois, image_meta] + feature_maps)
x = KL.TimeDistributed(KL.Conv2D(fc_layers_size, (pool_size, pool_size), padding="valid"),
                           name="mrcnn_class_conv1")(x)

where rois is generated by ProposalLayer from Mask_RCNN module (https://github.com/matterport/Mask_RCNN)

ProposalLayer inherits from the class keras.layers.Layer. In Keras the overwritten function compute_output_shape is called, but not in the tf.keras.

Additionally, the rois created in the tf.keras does not have the parameter _keras_shape that is most likely needed in the function TimeDistributed to create a proper shape.