Repository: tensorflow
Issue: Unexpected output shape when trying to convert to Frozen Graph using convert_variables_to_constants_v2
Labels: stat:awaiting tensorflower,type:bug,comp:core,2.6.0

<em>Please make sure that this is a bug. As per our
[GitHub Policy](https://github.com/tensorflow/tensorflow/blob/master/ISSUES.md),
we only address code/doc bugs, performance issues, feature requests and
build/installation issues on GitHub. tag:bug_template</em>

**System information**
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Debian
- Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: NA
- TensorFlow installed from (source or binary):  Using Deep Learning VM Instance on GCP. Not sure.
- TensorFlow version (use command below): 2.2.0
- Python version: 3.7.0
- Bazel version (if compiling from source): NA
- GCC/Compiler version (if compiling from source): NA
- CUDA/cuDNN version: 11.0
- GPU model and memory: Tesla T4 16GB

You can collect some of this information using our environment capture
[script](https://github.com/tensorflow/tensorflow/tree/master/tools/tf_env_collect.sh)
You can also obtain the TensorFlow version with:
1. TF 1.0: `python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"`
2. TF 2.0: `python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"`


**Describe the current behavior**

I am using this [answer](https://stackoverflow.com/a/61004639/5052482) by TF_Support on StackOverflow to convert a model to to a frozen graph. However, the output shape that I am getting for my custom model is is very different (weird) from the one that's there in my model. More information in the code snippet below.

Provide a reproducible test case that is the bare minimum necessary to generate
the problem. If possible, please share a link to Colab/Jupyter/any notebook.

```
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

def blah():
    jumps = tf.keras.layers.Input(shape=(269, 8))
    loc_a = tf.keras.layers.Input(shape=(32))
    loc_b = tf.keras.layers.Input(shape=(256))
    loc_c = tf.keras.layers.Input(shape=(2))
    def conv_block(x):
        x = tf.keras.layers.Conv1D(filters=32, kernel_size=1, padding='same')(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Activation('linear')(x)
        return x
    
    x = jumps
    for i in range(1):
        x = conv_block(x)
    _loc_a = tf.keras.layers.Dense(269)(loc_a)
    _loc_b = tf.keras.layers.Dense(269)(loc_b)
    _loc_c = tf.keras.layers.Dense(269)(loc_c)
    _loc = tf.keras.layers.Concatenate()([_loc_a, _loc_b, _loc_c])
    _loc = tf.keras.layers.Dense(269)(_loc)
    _loc = tf.keras.layers.Reshape((269, 1))(_loc)
    
    x = tf.keras.layers.Concatenate()([_loc, x])
    x= tf.keras.layers.Dense(9)(x)
#     x= tf.keras.layers.Dense(9)(x)
#     x = tf.keras.layers.Flatten()(x)
    
    return tf.keras.Model(inputs=[jumps, loc_a, loc_b, loc_c], outputs=[x])

model = blah()

model.summary()

Model: "model"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_2 (InputLayer)            [(None, 32)]         0                                            
__________________________________________________________________________________________________
input_3 (InputLayer)            [(None, 256)]        0                                            
__________________________________________________________________________________________________
input_4 (InputLayer)            [(None, 2)]          0                                            
__________________________________________________________________________________________________
dense (Dense)                   (None, 269)          8877        input_2[0][0]                    
__________________________________________________________________________________________________
dense_1 (Dense)                 (None, 269)          69133       input_3[0][0]                    
__________________________________________________________________________________________________
dense_2 (Dense)                 (None, 269)          807         input_4[0][0]                    
__________________________________________________________________________________________________
input_1 (InputLayer)            [(None, 269, 8)]     0                                            
__________________________________________________________________________________________________
concatenate (Concatenate)       (None, 807)          0           dense[0][0]                      
                                                                 dense_1[0][0]                    
                                                                 dense_2[0][0]                    
__________________________________________________________________________________________________
conv1d (Conv1D)                 (None, 269, 32)      288         input_1[0][0]                    
__________________________________________________________________________________________________
dense_3 (Dense)                 (None, 269)          217352      concatenate[0][0]                
__________________________________________________________________________________________________
batch_normalization (BatchNorma (None, 269, 32)      128         conv1d[0][0]                     
__________________________________________________________________________________________________
reshape (Reshape)               (None, 269, 1)       0           dense_3[0][0]                    
__________________________________________________________________________________________________
activation (Activation)         (None, 269, 32)      0           batch_normalization[0][0]        
__________________________________________________________________________________________________
concatenate_1 (Concatenate)     (None, 269, 33)      0           reshape[0][0]                    
                                                                 activation[0][0]                 
__________________________________________________________________________________________________
dense_4 (Dense)                 (None, 269, 9)       306         concatenate_1[0][0]              
==================================================================================================
Total params: 296,891
Trainable params: 296,827
Non-trainable params: 64

# convert model to a frozen graph

full_model = tf.function(lambda x: model(x))

full_model = full_model.get_concrete_function([tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype, name="input_0"),
                                              tf.TensorSpec(model.inputs[1].shape, model.inputs[0].dtype, name="input_1"),
                                              tf.TensorSpec(model.inputs[2].shape, model.inputs[0].dtype, name="input_2"),
                                              tf.TensorSpec(model.inputs[3].shape, model.inputs[0].dtype, name="input_3")])

frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def(add_shapes=True)
layers = [op.name for op in frozen_func.graph.get_operations()]
print("-" * 50)
# print("Frozen model layers: ")
# for layer in layers:
#     print(layer)
print("-" * 50)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)
# Save frozen graph from frozen ConcreteFunction to hard drive
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                  logdir="./frozen_models",
                  name="frozen_graph.pb",
                  as_text=False)

Frozen model inputs: 
[<tf.Tensor 'input_0:0' shape=(None, 269, 8) dtype=float32>, <tf.Tensor 'input_1:0' shape=(None, 32) dtype=float32>, <tf.Tensor 'input_2:0' shape=(None, 256) dtype=float32>, <tf.Tensor 'input_3:0' shape=(None, 2) dtype=float32>]
Frozen model outputs: 
[<tf.Tensor 'Identity:0' shape=(None, None, 9) dtype=float32>]
'./frozen_models/frozen_graph.pb'

```

The output shape of the Frozen Model should match the output shape of the model created. However it does not. 

The output shape of the model should be `(None, 269, 9)` whereas the output shape of the Frozen Graph is `(None, None, 9)`

**Other info / logs** Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.