I found that in new tf there are some problems with dynamic shapes. New tf use keras.layer and it impossible to make custom with dynamic shape input, the same in the plain code, tf.shape not return what it should.

Did you saw this problems?