How to properly define a model/layer which includes shape operation in TensorFlow?

Question

I encounter few different layout failed issues when my layers require shape operation like reshape, get input\'s shape. I carefully wrapped those ops with tensorflow ops lik