# Buggy code
def _balanced_data(self, data: Data, batch_size: int, shuffle: bool) -> Data:
    """Mix model data to account for class imbalance.

    This batching strategy puts rare classes approximately in every other batch,
    by repeating them. Mimics stratified batching, but also takes into account
    that more populated classes should appear more often.
    """

    self._check_label_key()

    # skip balancing if labels are token based
    if self.label_key is None or data[self.label_key][0][0].size > 1:
        return data

    label_ids = self._create_label_ids(data[self.label_key][0])

    unique_label_ids, counts_label_ids = np.unique(
        label_ids, return_counts=True, axis=0
    )
    num_label_ids = len(unique_label_ids)

    # group data points by their label
    # need to call every time, so that the data is shuffled inside each class
    data_by_label = self._split_by_label_ids(data, label_ids, unique_label_ids)

    # running index inside each data grouped by labels
    data_idx = [0] * num_label_ids
    # number of cycles each label was passed
    num_data_cycles = [0] * num_label_ids
    # if a label was skipped in current batch
    skipped = [False] * num_label_ids

    new_data = defaultdict(list)

    while min(num_data_cycles) == 0:
        if shuffle:
            indices_of_labels = np.random.permutation(num_label_ids)
        else:
            indices_of_labels = range(num_label_ids)

        for index in indices_of_labels:
            if num_data_cycles[index] > 0 and not skipped[index]:
                skipped[index] = True
                continue
            else:
                skipped[index] = False

            index_batch_size = (
                int(counts_label_ids[index] / self.num_examples * batch_size) + 1
            )

            for k, values in data_by_label[index].items():
                for i, v in enumerate(values):
                    if len(new_data[k]) < i + 1:
                        new_data[k].append([])
                    new_data[k][i].append(
                        v[data_idx[index] : data_idx[index] + index_batch_size]
                    )

            data_idx[index] += index_batch_size
            if data_idx[index] >= counts_label_ids[index]:
                num_data_cycles[index] += 1
                data_idx[index] = 0

            if min(num_data_cycles) > 0:
                break

    final_data = defaultdict(list)
    for k, values in new_data.items():
        for v in values:
            final_data[k].append(np.concatenate(np.array(v)))

    return final_data


# Traceback (ValueError)
Traceback (most recent call last):
File "/Users/tabergma/Repositories/rasa/rasa/__main__.py", line 119, in <module>
main()
File "/Users/tabergma/Repositories/rasa/rasa/__main__.py", line 108, in main
cmdline_arguments.func(cmdline_arguments)
File "/Users/tabergma/Repositories/rasa/rasa/cli/train.py", line 142, in train_nlu
additional_arguments=extract_nlu_additional_arguments(args),
File "/Users/tabergma/Repositories/rasa/rasa/train.py", line 431, in train_nlu
additional_arguments,
File "uvloop/loop.pyx", line 1456, in uvloop.loop.Loop.run_until_complete
File "/Users/tabergma/Repositories/rasa/rasa/train.py", line 472, in _train_nlu_async
additional_arguments=additional_arguments,
File "/Users/tabergma/Repositories/rasa/rasa/train.py", line 506, in _train_nlu_with_validated_data
**additional_arguments,
File "/Users/tabergma/Repositories/rasa/rasa/nlu/train.py", line 90, in train
interpreter = trainer.train(training_data, **kwargs)
File "/Users/tabergma/Repositories/rasa/rasa/nlu/model.py", line 191, in train
updates = component.train(working_data, self.config, **context)
File "/Users/tabergma/Repositories/rasa/rasa/nlu/classifiers/diet_classifier.py", line 716, in train
config=self.component_config,
File "/Users/tabergma/Repositories/rasa/rasa/nlu/classifiers/diet_classifier.py", line 1032, in __init__
self._check_data()
File "/Users/tabergma/Repositories/rasa/rasa/nlu/classifiers/diet_classifier.py", line 1099, in _check_data
"If hidden layer weights are shared, data signatures "
ValueError: If hidden layer weights are shared, data signatures for text_features and label_features must coincide.

# Fixed code
def _balanced_data(self, data: Data, batch_size: int, shuffle: bool) -> Data:
    """Mix model data to account for class imbalance.

    This batching strategy puts rare classes approximately in every other batch,
    by repeating them. Mimics stratified batching, but also takes into account
    that more populated classes should appear more often.
    """

    self._check_label_key()

    # skip balancing if labels are token based
    if self.label_key is None or data[self.label_key][0][0].size > 1:
        return data

    label_ids = self._create_label_ids(data[self.label_key][0])

    unique_label_ids, counts_label_ids = np.unique(
        label_ids, return_counts=True, axis=0
    )
    num_label_ids = len(unique_label_ids)

    # group data points by their label
    # need to call every time, so that the data is shuffled inside each class
    data_by_label = self._split_by_label_ids(data, label_ids, unique_label_ids)

    # running index inside each data grouped by labels
    data_idx = [0] * num_label_ids
    # number of cycles each label was passed
    num_data_cycles = [0] * num_label_ids
    # if a label was skipped in current batch
    skipped = [False] * num_label_ids

    new_data = defaultdict(list)

    while min(num_data_cycles) == 0:
        if shuffle:
            indices_of_labels = np.random.permutation(num_label_ids)
        else:
            indices_of_labels = range(num_label_ids)

        for index in indices_of_labels:
            if num_data_cycles[index] > 0 and not skipped[index]:
                skipped[index] = True
                continue

            skipped[index] = False

            index_batch_size = (
                int(counts_label_ids[index] / self.num_examples * batch_size) + 1
            )

            for k, values in data_by_label[index].items():
                for i, v in enumerate(values):
                    if len(new_data[k]) < i + 1:
                        new_data[k].append([])
                    new_data[k][i].append(
                        v[data_idx[index] : data_idx[index] + index_batch_size]
                    )

            data_idx[index] += index_batch_size
            if data_idx[index] >= counts_label_ids[index]:
                num_data_cycles[index] += 1
                data_idx[index] = 0

            if min(num_data_cycles) > 0:
                break

    final_data = defaultdict(list)
    for k, values in new_data.items():
        for v in values:
            final_data[k].append(np.concatenate(np.array(v)))

    return final_data