def check_and_transform_label_format(
    labels: np.ndarray, nb_classes: int | None, return_one_hot: bool = True
) -> np.ndarray:
    """
    Check label format and transform to one-hot-encoded labels if necessary. Only supports single-output classification.

    :param labels: An array of integer labels of shape `(nb_samples,)`, `(nb_samples, 1)` or `(nb_samples, nb_classes)`.
    :param nb_classes: The number of classes, as an integer. If None the number of classes is determined automatically.
    :param return_one_hot: True if returning one-hot encoded labels, False if returning index labels.
    :return: Labels with shape `(nb_samples, nb_classes)` (one-hot) or `(nb_samples,)` (index).
    """
    labels_return = labels

    if nb_classes is not None and not isinstance(nb_classes, (int, np.integer)):
        raise TypeError("nb_classes that is not an integer is not supported")

    if len(labels.shape) == 2 and labels.shape[1] > 1:  # multi-class, one-hot encoded
        if not return_one_hot:
            labels_return = np.argmax(labels, axis=1)
            labels_return = np.expand_dims(labels_return, axis=1)
    elif len(labels.shape) == 2 and labels.shape[1] == 1:
        if nb_classes is None:
            nb_classes = int(np.max(labels) + 1)
        if nb_classes > 2:  # multi-class, index labels
            if return_one_hot:
                labels_return = to_categorical(labels, nb_classes)
            else:
                labels_return = np.expand_dims(labels, axis=1)
        elif nb_classes == 2:  # binary, index labels
            if return_one_hot:
                labels_return = to_categorical(labels, nb_classes)
        else:
            raise ValueError(
                "Shape of labels not recognised."
                "Please provide labels in shape (nb_samples,) or (nb_samples, "
                "nb_classes)"
            )
    elif len(labels.shape) == 1:  # index labels
        if return_one_hot:
            labels_return = to_categorical(labels, nb_classes)
        else:
            labels_return = np.expand_dims(labels, axis=1)
    else:
        raise ValueError(
            "Shape of labels not recognised."
            "Please provide labels in shape (nb_samples,) or (nb_samples, "
            "nb_classes)"
        )

    return labels_return

    Check label format and transform to one-hot-encoded labels if necessary. Only supports single-output classification.

    :param labels: An array of integer labels of shape `(nb_samples,)`, `(nb_samples, 1)` or `(nb_samples, nb_classes)`.
    :param nb_classes: The number of classes, as an integer. If None the number of classes is determined automatically.
    :param return_one_hot: True if returning one-hot encoded labels, False if returning index labels.
    :return: Labels with shape `(nb_samples, nb_classes)` (one-hot) or `(nb_samples,)` (index).
    
art/utils.py