def get_io_shape(shape):
    """Gets Keras model I/O shape.

    This takes in the output of tf.keras.Model.Layer.input_shape or tf.keras.Model.Layer.output_shape.
    If the shape is already a tuple, simply returns the shape.
    If the shape is a list (as occasionally happens), it returns the shape tuple associated with this list.

    Args:
        shape: output of tf.keras.Model.Layer.input_shape or tf.keras.Model.Layer.output_shape
    """

    if type(shape) is list:
        if len(shape) != 1:
            raise ValueError("Invalid TF model I/O size")
        else:
            shape = shape[0]

    return shape