code: def get_state_shape_invariants(tensor):
    
    shape = tensor.shape.as_list()
    for i in range(1, len(shape) - 1):
      shape[i] = None
    return tf.TensorShape(shape)
docs: Returns the shape of the tensor but sets middle dims to None.