def get_shape(dataset):
    """
    Get the shape of the dataset.

    :param dataset: dataset to get the shape

    :returns: shape of the dataset
    """

    if isinstance(dataset, dict):
        transformed = {key: data.shape for key, data in dataset.items()}
    else:
        transformed = dataset.shape

    return transformed