def test_equal_class_distribution(self):
        """Verify equal class distribution in train and test sets."""
        (_, y_train), (_, y_test) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
        )
        _, counts_train = np.unique(y_train, return_counts=True)
        _, counts_test = np.unique(y_test, return_counts=True)

        self.assertTrue(
            np.all(counts_train == self.train_samples // self.num_classes)
        )
        self.assertTrue(
            np.all(counts_test == self.test_samples // self.num_classes)
        )
Verify equal class distribution in train and test sets.
keras/src/testing/test_utils_test.py