GitLab

Specify the axis explicitly when using `tf.squeueze` to avoid unknown result shape when the input contains dynamic shape dimensions.

PiperOrigin-RevId: 457858716