chore: import upstream snapshot with attribution

tests/conftest.py

@@ -0,0 +1,15 @@
+def pytest_collection_modifyitems(config, items):
+    """
+    Place the slowest tests at tne end of the test execution order if `order-tests` flag is set.
+    Adapted from https://stackoverflow.com/a/54254338/8505509
+    """
+    if config.option.order_tests:
+        items.sort(key=lambda item: 2 if item.get_closest_marker("slow") else 1)
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--order-tests",
+        action="store_true",
+        help="Flag to place slow test functions at the end of the test execution order",
+    )

tests/datalab/conftest.py

@@ -0,0 +1,128 @@
+import numpy as np
+import pandas as pd
+import pytest
+from datasets import load_dataset
+from datasets.arrow_dataset import Dataset
+from PIL import Image
+from sklearn.neighbors import NearestNeighbors
+
+from cleanlab.datalab.datalab import Datalab
+
+SEED = 42
+LABEL_NAME = "star"
+
+
+@pytest.fixture
+def dataset():
+    data_dict = {
+        "id": [
+            "7bd227d9-afc9-11e6-aba1-c4b301cdf627",
+            "7bd22905-afc9-11e6-a5dc-c4b301cdf627",
+            "7bd2299c-afc9-11e6-85d6-c4b301cdf627",
+            "7bd22a26-afc9-11e6-9309-c4b301cdf627",
+            "7bd22aba-afc9-11e6-8293-c4b301cdf627",
+        ],
+        "package_name": [
+            "com.mantz_it.rfanalyzer",
+            "com.mantz_it.rfanalyzer",
+            "com.mantz_it.rfanalyzer",
+            "com.mantz_it.rfanalyzer",
+            "com.mantz_it.rfanalyzer",
+        ],
+        "review": [
+            "Great app! The new version now works on my Bravia Android TV which is great as it's right by my rooftop aerial cable. The scan feature would be useful...any ETA on when this will be available? Also the option to import a list of bookmarks e.g. from a simple properties file would be useful.",
+            "Great It's not fully optimised and has some issues with crashing but still a nice app  especially considering the price and it's open source.",
+            "Works on a Nexus 6p I'm still messing around with my hackrf but it works with my Nexus 6p  Trond usb-c to usb host adapter. Thanks!",
+            "The bandwidth seemed to be limited to maximum 2 MHz or so. I tried to increase the bandwidth but not possible. I purchased this is because one of the pictures in the advertisement showed the 2.4GHz band with around 10MHz or more bandwidth. Is it not possible to increase the bandwidth? If not  it is just the same performance as other free APPs.",
+            "Works well with my Hackrf Hopefully new updates will arrive for extra functions",
+        ],
+        "date": [
+            "October 12 2016",
+            "August 23 2016",
+            "August 04 2016",
+            "July 25 2016",
+            "July 22 2016",
+        ],
+        "star": [4, 4, 5, 3, 5],
+        "version_id": [1487, 1487, 1487, 1487, 1487],
+    }
+    return Dataset.from_dict(data_dict)
+
+
+@pytest.fixture
+def label_name():
+    return LABEL_NAME
+
+
+@pytest.fixture
+def lab(dataset, label_name):
+    return Datalab(data=dataset, label_name=label_name)
+
+
+@pytest.fixture
+def large_lab():
+    np.random.seed(SEED)
+    N = 100
+    K = 2
+    data = np.random.rand(N, 2)
+    labels = np.random.randint(0, K, size=N)
+    pred_probs = np.random.rand(N, K)
+    pred_probs /= pred_probs.sum(axis=1, keepdims=True)
+
+    lab = Datalab(
+        data={"features": data, "label": labels, "pred_probs": pred_probs}, label_name="label"
+    )
+    knn = NearestNeighbors(n_neighbors=25, metric="euclidean").fit(data)
+    knn_graph = knn.kneighbors_graph(mode="distance")
+    lab.info["statistics"]["unit_test_knn_graph"] = knn_graph
+    return lab
+
+
+@pytest.fixture
+def pred_probs(dataset):
+    np.random.seed(SEED)
+    return np.random.rand(len(dataset), 3)
+
+
+@pytest.fixture
+def custom_issue_manager():
+    from cleanlab.datalab.internal.issue_manager.issue_manager import IssueManager
+
+    class CustomIssueManager(IssueManager):
+        issue_name = "custom_issue"
+
+        def find_issues(self, custom_argument: int = 1, **_) -> None:
+            # Flag example as an issue if the custom argument equals its index
+            scores = [
+                abs(i - custom_argument) / (i + custom_argument)
+                for i in range(len(self.datalab.data))
+            ]
+            self.issues = pd.DataFrame(
+                {
+                    f"is_{self.issue_name}_issue": [
+                        i == custom_argument for i in range(len(self.datalab.data))
+                    ],
+                    self.issue_score_key: scores,
+                },
+            )
+            summary_score = np.mean(scores)
+            self.summary = self.make_summary(score=summary_score)
+
+    return CustomIssueManager
+
+
+def generate_image():
+    arr = np.random.randint(low=0, high=256, size=(300, 300, 3), dtype=np.uint8)
+    img = Image.fromarray(arr, mode="RGB")
+    return img
+
+
+@pytest.fixture
+def image_dataset():
+    data_path = "./tests/datalab/data"
+    dataset = load_dataset(
+        "imagefolder",
+        data_dir=data_path,
+        split="train",
+    )
+    return dataset

tests/datalab/datalab/end_to_end_tests/issue_types/test_null.py

@@ -0,0 +1,80 @@
+import numpy as np
+import pandas as pd
+import pytest
+from cleanlab import Datalab
+from cleanlab.datalab.internal.issue_manager.null import NullIssueManager
+
+
+@pytest.fixture
+def setup_test_environment():
+    """Setup the basic test environment with a Datalab instance and a NullIssueManager."""
+    data = {
+        "features": np.array([[1, 2, np.nan], [4, np.nan, 6], [np.nan, 8, 9]]),
+        "label": np.array([0, 1, 0]),
+    }
+    lab = Datalab(data=data, label_name="label")
+    issue_manager = NullIssueManager(lab)
+    return {"lab": lab, "issue_manager": issue_manager}
+
+
+class FeatureFactory:
+    """Factory class to create features in different formats."""
+
+    @staticmethod
+    def create_features(format_type="numpy"):
+        """Generate test features in the specified format."""
+        strategy = {
+            "numpy": FeatureFactory._create_numpy_features,
+            "pandas": FeatureFactory._create_pandas_features,
+        }
+        feature_strategy = strategy.get(format_type)
+        if feature_strategy is None:
+            raise ValueError(
+                f"Invalid format_type to test: {format_type}, must be one of {list(strategy.keys())}"
+            )
+        return feature_strategy()
+
+    @staticmethod
+    def _create_numpy_features():
+        """Generate features in numpy format."""
+        return np.array([[1, 2, np.nan], [4, np.nan, 6], [np.nan, 8, 9]])
+
+    @staticmethod
+    def _create_pandas_features():
+        """Generate features in pandas DataFrame format."""
+        features = FeatureFactory._create_numpy_features()
+        return pd.DataFrame(features, columns=["a", "b", "c"])
+
+
+@pytest.mark.parametrize("format_type", ["numpy", "pandas"])
+def test_null_issue_manager(setup_test_environment, format_type):
+    features = FeatureFactory.create_features(format_type)
+    null_issue_manager = setup_test_environment["issue_manager"]
+    null_issue_manager.find_issues(features=features)
+    assert null_issue_manager.info["most_common_issue"] == {
+        "pattern": "001",
+        "rows_affected": [0],
+        "count": 1,
+    }
+
+
+@pytest.mark.parametrize("format_type", ["numpy", "pandas"])
+def test_lab_find_issues(setup_test_environment, format_type):
+    features = FeatureFactory.create_features(format_type)
+    lab = setup_test_environment["lab"]
+    lab.find_issues(features=features, issue_types={"null": {}})
+
+    null_issues = lab.get_issues("null")
+    expected_null_issues = pd.DataFrame(
+        {
+            "is_null_issue": [False] * 3,
+            "null_score": [2 / 3] * 3,
+        }
+    )
+    pd.testing.assert_frame_equal(null_issues, expected_null_issues)
+
+    most_common_issue = lab.get_info("null")["most_common_issue"]
+    assert most_common_issue == {"pattern": "001", "rows_affected": [0], "count": 1}
+
+    column_impact = lab.get_info("null")["column_impact"]
+    np.testing.assert_array_equal(column_impact, [1 / 3] * 3)

tests/datalab/datalab/end_to_end_tests/test_all_identical_examples.py

@@ -0,0 +1,367 @@
+import numpy as np
+import pandas as pd
+import pytest
+from cleanlab import Datalab
+
+
+class TestAllIdenticalExamplesDataset:
+    """There are 4 types of datasets this class tests:
+
+    1. A dataset with all identical data points. They all have identical labels, EXCEPT THE LAST ONE.
+    2. A dataset with all identical data points, except for one extra/unique example. One of the identical examples has a noisy label. The unique example has the same label as the majority of the dataset.
+    3. A regression dataset with all identical data points. They all have identical targets, EXCEPT THE LAST ONE.
+    4. A regression dataset with all identical data points, except for one extra/unique example. One of the identical examples has a noisy target. The unique example has the same target as the majority of the dataset.
+
+    Each test goes through the same motions:
+
+    - Set up Datalab instance
+    - Find default issues from features and pred_probs/predictions
+    - Assert that the following issue dataframes look as expected for each dataset type.
+      - label issues
+      - outlier issues
+      - near-duplicate issues
+
+    There are a few more issue types that are tested in some cases, but they are less prone to be affected by future changes in the codebase.
+    - underperforming group issues (for classification)
+    - class imbalance issues (for classification)
+
+    These tests focus on the issue detection capabilities of Datalab, and not the quality of the predictions or the features.
+    """
+
+    @pytest.fixture
+    def dataset(self, request):
+        N, M = request.param
+        # Create the dataset with all identical examples
+        X = np.full((N, M), fill_value=np.random.rand(M))
+
+        # All labels for identical points should be the same
+        y = ["a"] * N
+        # One of the labels is different, so it should be flagged with a label issue
+        y[-1] = "b"
+
+        return {"X": X, "y": y}
+
+    @pytest.fixture
+    def dataset_with_one_unique_example(self, request):
+        N, M = request.param
+        # Create the dataset with all identical examples
+        X = np.full((N, M), fill_value=np.random.rand(M))
+
+        # Add one unique example to the dataset
+        X = np.vstack([X, np.random.rand(M)])
+
+        # All labels for identical points should be the same, let's make them all 0, along with the unique example
+        y = ["a"] * (N + 1)
+        # One of the N points has a noisy label, it should be flagged with a label issue
+        y[-2] = "b"
+
+        return {"X": X, "y": y}
+
+    @pytest.fixture
+    def regression_dataset(self, request):
+        N, M = request.param
+        # Create the dataset with all identical examples
+        X = np.full((N, M), fill_value=np.random.rand(M))
+
+        # All labels for identical points should be the same
+        y = np.full(N, fill_value=np.random.rand())
+        # Flip one of the targets to introduce a label issue
+        y[-1] += 10
+
+        return {"X": X, "y": y}
+
+    @pytest.fixture
+    def regression_dataset_with_one_unique_example(self, request):
+        N, M = request.param
+        # Create the dataset with all identical examples
+        X = np.full((N, M), fill_value=np.random.rand(M))
+
+        # All labels for identical points should be the same
+        y = np.full(N, fill_value=np.random.rand())
+        # Flip one of the targets to introduce a label issue
+        y[-1] += 10
+
+        # Add one unique example to the dataset, but it has the same target as the majority of the dataset
+        X = np.vstack([X, np.random.rand(M)])
+        y = np.append(y, [y[0]])
+
+        return {"X": X, "y": y}
+
+    @pytest.mark.parametrize(
+        "dataset",
+        [((N, M)) for N in [11, 20, 50, 100, 150] for M in [2, 3, 5, 10, 20]],
+        indirect=["dataset"],
+        ids=lambda x: f"N={x[0]}, M={x[1]}",
+    )
+    def test_issue_detection(self, dataset):
+        lab = Datalab(data=dataset, label_name="y")
+        N = len(dataset["y"])
+        pred_probs = np.full((N, 2), fill_value=[1.0, 0.0])
+        lab.find_issues(features=dataset["X"], pred_probs=pred_probs)
+
+        outlier_issues = lab.get_issues("outlier")
+        expected_outlier_issues = pd.DataFrame(
+            [{"is_outlier_issue": False, "outlier_score": 1.0}] * N
+        )
+        pd.testing.assert_frame_equal(outlier_issues, expected_outlier_issues)
+
+        near_duplicate_issues = lab.get_issues("near_duplicate")
+        expected_near_duplicate_issues = pd.DataFrame(
+            [
+                {
+                    "is_near_duplicate_issue": True,
+                    "near_duplicate_score": 0.0,
+                    "distance_to_nearest_neighbor": np.finfo(np.float64).epsneg,
+                }
+                for i in range(N)
+            ]
+        )
+        pd.testing.assert_frame_equal(
+            near_duplicate_issues.drop(columns="near_duplicate_sets"),
+            expected_near_duplicate_issues,
+            check_exact=False,
+            atol=5e-16,
+        )
+
+        label_issues = lab.get_issues("label")[["is_label_issue", "label_score"]]
+        expected_label_issues = pd.DataFrame(
+            [{"is_label_issue": False, "label_score": 1.0}] * (N - 1)
+            + [
+                {"is_label_issue": False, "label_score": 0.0}
+            ]  # The confident threshold for this examble is extremely low, but won't flag it as a label issue
+        )
+        pd.testing.assert_frame_equal(label_issues, expected_label_issues)
+
+        underperforming_group_issues = lab.get_issues("underperforming_group")
+        expected_underperforming_group_issues = pd.DataFrame(
+            [
+                {
+                    "is_underperforming_group_issue": False,
+                    "underperforming_group_score": 1.0,
+                }
+            ]
+            * N  # Only a single data point performs poorly, so it's not in any group
+        )
+        pd.testing.assert_frame_equal(
+            underperforming_group_issues, expected_underperforming_group_issues
+        )
+
+        if N > 20:
+            # Default threshold for class imbalance is 0.1*1/K, where K is the number of classes. So for 20 examples and 2 classes,
+            # the threshold is 5% (so 1 example out of 20 is NOT considered class imbalance)
+            class_imbalance_issues = lab.get_issues("class_imbalance")[
+                ["is_class_imbalance_issue", "class_imbalance_score"]
+            ]
+            expected_class_imbalance_issues = pd.DataFrame(
+                [
+                    {
+                        "is_class_imbalance_issue": False,
+                        "class_imbalance_score": 1.0,
+                    }
+                ]
+                * (N - 1)
+                + [
+                    {
+                        "is_class_imbalance_issue": True,
+                        "class_imbalance_score": 1 / N,
+                    }
+                ]
+            )
+            pd.testing.assert_frame_equal(class_imbalance_issues, expected_class_imbalance_issues)
+
+    @pytest.mark.parametrize(
+        "dataset_with_one_unique_example",
+        [((N, M)) for N in [11, 20, 50, 100, 150] for M in [2, 3, 5, 10, 20]],
+        indirect=["dataset_with_one_unique_example"],
+        ids=lambda x: f"N={x[0]}, M={x[1]}",
+    )
+    def test_issue_detection_with_one_unique_example(self, dataset_with_one_unique_example):
+        dataset = dataset_with_one_unique_example
+        N = len(dataset["y"])
+        lab = Datalab(data=dataset, label_name="y")
+        pred_probs = np.full((N, 2), fill_value=[1.0, 0.0])
+        lab.find_issues(features=dataset["X"], pred_probs=pred_probs)
+
+        outlier_issues = lab.get_issues("outlier")
+        expected_outlier_issues = pd.DataFrame(
+            [{"is_outlier_issue": False, "outlier_score": 1.0}] * (N - 1)
+            + [{"is_outlier_issue": True, "outlier_score": 0.0}]
+        )
+
+        pd.testing.assert_frame_equal(outlier_issues, expected_outlier_issues)
+
+        near_duplicate_issues = lab.get_issues("near_duplicate")
+        expected_near_duplicate_issues = pd.DataFrame(
+            [
+                {
+                    "is_near_duplicate_issue": True,
+                    "near_duplicate_score": 0.0,
+                }
+                for i in range(N - 1)
+            ]
+            + [
+                {
+                    "is_near_duplicate_issue": False,
+                    "near_duplicate_score": 1.0,
+                }
+            ]
+        )
+        pd.testing.assert_frame_equal(
+            near_duplicate_issues.drop(
+                columns=["near_duplicate_sets", "distance_to_nearest_neighbor"]
+            ),
+            expected_near_duplicate_issues,
+        )
+
+        label_issues = lab.get_issues("label")[["is_label_issue", "label_score"]]
+        expected_label_issues = pd.DataFrame(
+            [{"is_label_issue": False, "label_score": 1.0}] * (N - 2)
+            + [
+                {"is_label_issue": False, "label_score": 0.0}
+            ]  # The confident threshold for this examble is extremely low, but won't flag it as a label issue
+            + [{"is_label_issue": False, "label_score": 1.0}]
+        )
+        pd.testing.assert_frame_equal(label_issues, expected_label_issues)
+
+        underperforming_group_issues = lab.get_issues("underperforming_group")
+        expected_underperforming_group_issues = pd.DataFrame(
+            [
+                {
+                    "is_underperforming_group_issue": False,
+                    "underperforming_group_score": 1.0,
+                }
+            ]
+            * N  # Only a single data point performs poorly, so it's not in any group
+        )
+        pd.testing.assert_frame_equal(
+            underperforming_group_issues, expected_underperforming_group_issues
+        )
+
+        if N > 20:
+            # Default threshold for class imbalance is 0.1*1/K, where K is the number of classes. So for 20 examples and 2 classes,
+            # the threshold is 5% (so 1 example out of 20 is NOT considered class imbalance)
+            class_imbalance_issues = lab.get_issues("class_imbalance")[
+                ["is_class_imbalance_issue", "class_imbalance_score"]
+            ]
+            expected_class_imbalance_issues = pd.DataFrame(
+                [
+                    {
+                        "is_class_imbalance_issue": False,
+                        "class_imbalance_score": 1.0,
+                    }
+                ]
+                * (N - 2)
+                + [
+                    {
+                        "is_class_imbalance_issue": True,
+                        "class_imbalance_score": 1 / N,
+                    }
+                ]
+                + [
+                    {
+                        "is_class_imbalance_issue": False,
+                        "class_imbalance_score": 1.0,
+                    }
+                ]
+            )
+            pd.testing.assert_frame_equal(class_imbalance_issues, expected_class_imbalance_issues)
+
+    @pytest.mark.parametrize(
+        "regression_dataset",
+        [((N, M)) for N in [11, 20, 50, 100] for M in [3, 4, 6, 8, 10]],
+        indirect=["regression_dataset"],
+        ids=lambda x: f"N={x[0]}, M={x[1]}",
+    )
+    def test_regression_issue_detection(self, regression_dataset):
+        lab = Datalab(data=regression_dataset, label_name="y", task="regression")
+        predictions = np.full(len(y := regression_dataset["y"]), fill_value=y[0])
+
+        lab.find_issues(pred_probs=predictions, features=regression_dataset["X"])
+
+        outlier_issues = lab.get_issues("outlier")
+        expected_outlier_issues = pd.DataFrame(
+            [{"is_outlier_issue": False, "outlier_score": 1.0}] * len(outlier_issues)
+        )
+        pd.testing.assert_frame_equal(outlier_issues, expected_outlier_issues)
+
+        near_duplicate_issues = lab.get_issues("near_duplicate")
+        expected_near_duplicate_issues = pd.DataFrame(
+            [
+                {
+                    "is_near_duplicate_issue": True,
+                    "near_duplicate_score": 0.0,
+                    "distance_to_nearest_neighbor": np.finfo(np.float64).epsneg,
+                }
+                for i in range(len(near_duplicate_issues))
+            ]
+        )
+        pd.testing.assert_frame_equal(
+            near_duplicate_issues.drop(columns="near_duplicate_sets"),
+            expected_near_duplicate_issues,
+            check_exact=False,
+            atol=5e-16,
+        )
+
+        label_issues = lab.get_issues("label")[["is_label_issue", "label_score"]]
+        expected_label_issues = pd.DataFrame(
+            [{"is_label_issue": False, "label_score": 1.0}] * (len(label_issues) - 1)
+            + [
+                {"is_label_issue": True, "label_score": 0.0}
+            ]  # Expect last example to be flagged as a label issue
+        )
+        pd.testing.assert_frame_equal(label_issues, expected_label_issues)
+
+    @pytest.mark.parametrize(
+        "regression_dataset_with_one_unique_example",
+        [((N, M)) for N in [11, 20, 50, 100] for M in [3, 4, 6, 8, 10]],
+        indirect=["regression_dataset_with_one_unique_example"],
+        ids=lambda x: f"N={x[0]}, M={x[1]}",
+    )
+    def test_regression_issue_detection_with_one_unique_example(
+        self, regression_dataset_with_one_unique_example
+    ):
+        dataset = regression_dataset_with_one_unique_example
+        lab = Datalab(data=dataset, label_name="y", task="regression")
+        predictions = np.full(len(y := dataset["y"]), fill_value=y[0])
+
+        lab.find_issues(pred_probs=predictions, features=dataset["X"])
+
+        outlier_issues = lab.get_issues("outlier")
+        expected_outlier_issues = pd.DataFrame(
+            [{"is_outlier_issue": False, "outlier_score": 1.0}] * (len(outlier_issues) - 1)
+            + [{"is_outlier_issue": True, "outlier_score": 0.0}]
+        )
+        pd.testing.assert_frame_equal(outlier_issues, expected_outlier_issues)
+
+        near_duplicate_issues = lab.get_issues("near_duplicate")
+        expected_near_duplicate_issues = pd.DataFrame(
+            [
+                {
+                    "is_near_duplicate_issue": True,
+                    "near_duplicate_score": 0.0,
+                }
+                for i in range(len(near_duplicate_issues) - 1)
+            ]
+            + [
+                {
+                    "is_near_duplicate_issue": False,
+                    "near_duplicate_score": 1.0,
+                }
+            ]
+        )
+        pd.testing.assert_frame_equal(
+            near_duplicate_issues.drop(
+                columns=["near_duplicate_sets", "distance_to_nearest_neighbor"]
+            ),
+            expected_near_duplicate_issues,
+        )
+
+        label_issues = lab.get_issues("label")[["is_label_issue", "label_score"]]
+        expected_label_issues = pd.DataFrame(
+            [{"is_label_issue": False, "label_score": 1.0}] * (len(label_issues) - 2)
+            + [
+                {"is_label_issue": True, "label_score": 0.0}
+            ]  # Expect second to last example to be flagged as a label issue
+            + [{"is_label_issue": False, "label_score": 1.0}]
+        )
+        pd.testing.assert_frame_equal(label_issues, expected_label_issues)

tests/datalab/datalab/test_multilabel.py

@@ -0,0 +1,162 @@
+import numpy as np
+import pytest
+from sklearn.neighbors import NearestNeighbors
+
+
+from cleanlab.datalab.datalab import Datalab
+from cleanlab.internal.multilabel_utils import int2onehot
+
+
+from cleanlab.benchmarking.noise_generation import (
+    generate_noise_matrix_from_trace,
+    generate_noisy_labels,
+)
+
+
+class TestDatalabForMultilabelClassification:
+    @pytest.fixture
+    def multilabel_data(
+        self,
+        means=[[-5, 3.5], [0, 2], [-3, 6]],
+        covs=[[[3, -1.5], [-1.5, 1]], [[5, -1.5], [-1.5, 1]], [[3, -1.5], [-1.5, 1]]],
+        boxes_coordinates=[[-3.5, 0, -1.5, 1.7], [-1, 3, 2, 4], [-5, 2, -3, 4], [-3, 2, -1, 4]],
+        box_multilabels=[[0, 1], [1, 2], [0, 2], [0, 1, 2]],
+        sizes=[100, 80, 100],
+        avg_trace=0.9,
+        seed=5,
+    ):
+        np.random.seed(seed=seed)
+        n = sum(sizes)
+        num_classes = len(means)
+        m = num_classes + len(
+            box_multilabels
+        )  # number of classes by treating each multilabel as 1 unique label
+        local_data = []
+        labels = []
+        for i in range(0, len(means)):
+            local_data.append(
+                np.random.multivariate_normal(mean=means[i], cov=covs[i], size=sizes[i])
+            )
+            labels += [[i]] * sizes[i]
+
+        def make_multi(X, Y, bx1, by1, bx2, by2, label_list):
+            ll = np.array([bx1, by1])  # lower-left
+            ur = np.array([bx2, by2])  # upper-right
+
+            inidx = np.all(np.logical_and(X.tolist() >= ll, X.tolist() <= ur), axis=1)
+            for i in range(0, len(Y)):
+                if inidx[i]:
+                    Y[i] = label_list
+            return Y
+
+        X_train = np.vstack(local_data)
+
+        for i in range(0, len(box_multilabels)):
+            bx1, by1, bx2, by2 = boxes_coordinates[i]
+            multi_label = box_multilabels[i]
+            labels = make_multi(X_train, labels, bx1, by1, bx2, by2, multi_label)
+
+        d = {}
+        for i in labels:
+            if str(i) not in d:
+                d[str(i)] = len(d)
+        inv_d = {v: k for k, v in d.items()}
+        labels_idx = [d[str(i)] for i in labels]
+        py = np.bincount(labels_idx) / float(len(labels_idx))
+        noise_matrix = generate_noise_matrix_from_trace(
+            m,
+            trace=avg_trace * m,
+            py=py,
+            valid_noise_matrix=True,
+            seed=seed,
+        )
+        noisy_labels_idx = generate_noisy_labels(labels_idx, noise_matrix)
+        noisy_labels = [eval(inv_d[i]) for i in noisy_labels_idx]
+        error_idx = np.where(labels_idx != noisy_labels_idx)[0]
+        pred_probs = np.full((n, num_classes), fill_value=0.1)
+        labels_onehot = int2onehot(labels, K=num_classes)
+        pred_probs[labels_onehot == 1] = 0.9
+
+        knn_graph = (
+            NearestNeighbors(n_neighbors=15, metric="euclidean")
+            .fit(X_train)
+            .kneighbors_graph(mode="distance")
+        )
+        return {
+            "X": X_train,
+            "true_y": labels,
+            "y": noisy_labels,
+            "error_idx": error_idx,
+            "pred_probs": pred_probs,
+            "knn_graph": knn_graph,
+        }
+
+    @pytest.fixture
+    def lab(self, multilabel_data):
+        X, y = multilabel_data["X"], multilabel_data["y"]
+        data = {"X": X, "y": y}
+        lab = Datalab(data=data, label_name="y", task="multilabel")
+        return lab
+
+    def test_available_issue_types(self, lab):
+        assert set(lab.list_default_issue_types()) == set(
+            ["label", "near_duplicate", "non_iid", "outlier", "null"]
+        )
+        assert set(lab.list_possible_issue_types()) == set(
+            ["label", "near_duplicate", "non_iid", "outlier", "null", "data_valuation"]
+        )
+
+    @pytest.mark.parametrize(
+        "argument_name, data_key",
+        [
+            ("pred_probs", "pred_probs"),
+            # TODO: Add support for finding multilabel issues from features
+            # ("features", "X"),
+            # TODO: Add support for finding multilabel issues from knn_graph
+            # ("knn_graph", "knn_graph"),
+        ],
+        ids=[
+            "pred_probs only",
+            # "features only",
+            # "knn_graph only",
+        ],
+    )
+    def test_find_label_issues(self, lab, multilabel_data, argument_name, data_key):
+        """Test that the multilabel classification issue checks finds at least 90% of the
+        label issues."""
+        input_dict = {argument_name: multilabel_data[data_key]}
+        issue_types = {"label": {}}
+        lab.find_issues(**input_dict, issue_types=issue_types)
+        lab.report()
+
+        issues = lab.get_issues("label")
+        issue_ids = issues.query("is_label_issue").index
+        expected_issue_ids = multilabel_data["error_idx"]
+
+        # jaccard similarity
+        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
+        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
+        assert float(intersection) / union >= 0.9
+
+    @pytest.mark.parametrize(
+        "argument_name, data_key, expected_issue_types_in_summary",
+        [
+            ("pred_probs", "pred_probs", ["label"]),
+            ("features", "X", ["near_duplicate", "non_iid", "outlier", "null"]),
+            ("knn_graph", "knn_graph", ["near_duplicate", "non_iid", "outlier"]),
+        ],
+        ids=[
+            "pred_probs only",
+            "features only",
+            "knn_graph only",
+        ],
+    )
+    def test_find_issues_defaults(
+        self, lab, multilabel_data, argument_name, data_key, expected_issue_types_in_summary
+    ):
+        """Test that the multilabel classification issue checks for various issues by default."""
+        input_dict = {argument_name: multilabel_data[data_key]}
+        lab.find_issues(**input_dict)
+        issue_summary = lab.get_issue_summary()
+        assert issue_summary["num_issues"].sum() > 0
+        assert set(issue_summary["issue_type"].values) == set(expected_issue_types_in_summary)

tests/datalab/datalab/test_regression.py

@@ -0,0 +1,635 @@
+import numpy as np
+import pandas as pd
+import pytest
+from sklearn.neighbors import NearestNeighbors
+from cleanlab.datalab.datalab import Datalab
+
+
+SEED = 42
+
+
+class TestDatalabForRegression:
+    true_columns = [
+        "true_label_error_idx",
+        "true_outlier_idx",
+        "true_duplicate_idx",
+        "true_non_iid_idx",
+        "true_underperforming_group_idx",
+    ]
+
+    @pytest.fixture
+    def regression_data(self, num_examples=400, num_features=3, error_frac=0.025, error_noise=0.5):
+        np.random.seed(SEED)
+        X = np.random.random(size=(num_examples, num_features))
+        coefficients = np.random.uniform(-5, 5, size=num_features)
+        coefficients[2] *= 0.005  # make the third feature less important
+
+        num_errors = int(num_examples * error_frac)
+        label_error_idx = np.random.choice(num_examples, num_errors)
+
+        # Move the examples with label errors closer to the main mass of examples
+        direction_towards_center = X[label_error_idx] - 0.5
+        direction_towards_center /= np.linalg.norm(direction_towards_center, axis=1)[:, None]
+        X[label_error_idx] -= direction_towards_center * 0.2
+
+        # Add near-duplicates
+        n_duplicates = 5
+        X = np.vstack([X, X[:n_duplicates] + 1e-6 * np.random.randn(n_duplicates, num_features)])
+        true_duplicate_ids = np.hstack(
+            [np.arange(n_duplicates), np.arange(num_examples, num_examples + n_duplicates)]
+        )
+
+        # Add outliers
+        n_outliers = 2
+        X_outliers = np.random.normal(
+            loc=[0.5, 0.5, 1.2], scale=0.05, size=(n_outliers, num_features)
+        )
+        X = np.vstack([X, X_outliers])
+        true_outlier_ids = len(X) - n_outliers + np.arange(n_outliers)
+
+        # Non-iid examples
+        # n_non_iid = 15
+        # # (linearly increasing features with some Gaussian noise)
+        # X_non_iid = np.linspace(0, 0.005, n_non_iid)[:, None] + np.random.normal(
+        #     loc=[0.5, 0.5, 0.1], scale=0.005, size=(n_non_iid, num_features)
+        # )
+        # X = np.vstack([X, X_non_iid])
+        # true_non_iid_ids = len(X) - n_non_iid + np.arange(n_non_iid)
+
+        # Underperforming group
+        n_underperforming = 10
+        # A linearly spaced blob with several targets that are too noisy
+
+        X_underperforming = np.array(
+            [
+                np.linspace(0, 0.025, n_underperforming),
+                np.linspace(0, 0.010, n_underperforming),
+                [0] * n_underperforming,
+            ]
+        ).T + np.random.normal(
+            loc=[-0.05, 1.01, 1.01], scale=0.0005, size=(n_underperforming, num_features)
+        )
+        X = np.vstack([X, X_underperforming])
+        true_underperforming_group_ids = len(X) - n_underperforming + np.arange(n_underperforming)
+
+        # Use the underperforming group as the non-iid group
+        true_non_iid_ids = np.copy(true_underperforming_group_ids)
+
+        true_y = np.dot(X, coefficients)
+
+        # add extra noisy examples
+
+        label_noise = np.clip(
+            np.random.normal(loc=error_noise, scale=error_noise / 4, size=num_errors),
+            1.5 * error_noise,
+            4 * error_noise,
+        ) * np.random.choice([-1, 1], size=num_errors)
+        y = true_y.copy()
+        y[label_error_idx] += label_noise
+        error_idx = np.argsort(abs(y - true_y))[-num_errors:]  # get the noisiest examples idx
+
+        # Add some noise to the underperforming group
+        y[true_underperforming_group_ids[-4:]] += np.random.normal(loc=0, scale=1e-6, size=4)
+
+        # Ensure that the MSE is not too large
+        from sklearn.linear_model import LinearRegression
+
+        linear_model = LinearRegression()
+
+        # Validate that the label noise affects the MSE for a linear model
+        from sklearn.metrics import mean_squared_error
+
+        y_pred_true = linear_model.fit(X, true_y).predict(X)
+        mse_true = mean_squared_error(true_y, y_pred_true)
+        assert float(mse_true) < 1e-10
+
+        y_pred = linear_model.fit(X, y).predict(X)
+        mse = mean_squared_error(y, y_pred)
+        assert 1e-3 < float(mse) < 5e-2
+
+        knn_graph = (
+            NearestNeighbors(n_neighbors=10, metric="euclidean")
+            .fit(X)
+            .kneighbors_graph(mode="distance")
+        )
+        return {
+            "X": X,
+            "y": y,
+            "true_y": true_y,
+            "knn_graph": knn_graph,
+            "error_idx": error_idx,
+            "duplicate_ids": true_duplicate_ids,
+            "outlier_ids": true_outlier_ids,
+            "non_iid_ids": true_non_iid_ids,
+            "underperforming_group_ids": true_underperforming_group_ids,
+        }
+
+    @pytest.fixture
+    def lab(self, regression_data):
+        X, y = regression_data["X"], regression_data["y"]
+        test_df = pd.DataFrame(X, columns=["c1", "c2", "c3"])
+        test_df["y"] = y
+        lab = Datalab(data=test_df, label_name="y", task="regression")
+        return lab
+
+    def test_available_issue_types(self, lab):
+        assert set(lab.list_default_issue_types()) == set(
+            ["label", "outlier", "near_duplicate", "non_iid", "null"]
+        )
+        assert set(lab.list_possible_issue_types()) == set(
+            ["label", "outlier", "near_duplicate", "non_iid", "null", "data_valuation"]
+        )
+
+    def test_regression_with_features_finds_label_issues(self, lab, regression_data):
+        """Test that the regression issue checks finds 40 label issues, based on the
+        numerical features."""
+        X = regression_data["X"]
+        issue_types = {"label": {"clean_learning_kwargs": {"seed": SEED}, "fine_search_size": 10}}
+        lab.find_issues(features=X, issue_types=issue_types)
+        lab.report()
+
+        issues = lab.get_issues("label")
+        issue_ids = issues.query("is_label_issue").index
+        expected_issue_ids = regression_data["error_idx"]
+
+        # jaccard similarity
+        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
+        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
+        expected_jaccard_similarity_bound = 0.7
+        has_high_jaccard = float(intersection) / union > expected_jaccard_similarity_bound
+        if not has_high_jaccard:
+            # Then check if MAP@len(expected_issue_ids) is high for the scores
+            def average_precision_at_k(y_true, y_pred, k) -> float:
+                """Compute average precision at k."""
+                y_true = np.asarray(y_true)
+                y_pred = np.asarray(y_pred)
+                sort_indices = np.argsort(y_pred)[::-1]
+                y_true = y_true[sort_indices]
+                return float(np.mean(y_true[:k]))
+
+            ap_at_k = average_precision_at_k(
+                # y_true are the expected issues
+                y_true=np.isin(np.arange(len(issues)), expected_issue_ids),
+                # y_pred are the scores
+                y_pred=1 - issues["label_score"],
+                k=len(expected_issue_ids),
+            )
+            ap_at_k_bound = 0.9
+            ap_at_k_is_high_enough = ap_at_k > ap_at_k_bound
+            print(f"AP@{len(expected_issue_ids)}: {ap_at_k}")
+            assert (
+                ap_at_k_is_high_enough
+            ), f"AP@{len(expected_issue_ids)} should be > {ap_at_k_bound}, got {ap_at_k}"
+
+        # FPR
+        fpr = len(list(set(issue_ids).difference(set(expected_issue_ids)))) / len(issue_ids)
+        expected_fpr_bound = 0.2
+        has_low_fpr = fpr < expected_fpr_bound
+
+        if not (has_high_jaccard and has_low_fpr):
+            ids_of_union = list(set(issue_ids).union(set(expected_issue_ids)))
+            _issues = issues.loc[ids_of_union]
+            ids_keys = [
+                "error_idx",
+                "outlier_ids",
+                "duplicate_ids",
+                "non_iid_ids",
+                "underperforming_group_ids",
+            ]
+            for new_col, ids_key in zip(self.true_columns, ids_keys):
+                _issues[new_col] = _issues.index.isin(regression_data[ids_key])
+            # Display dataframe
+            print("Expected label issue ids", expected_issue_ids, "\n")
+            _label_issues = _issues.query("is_label_issue")
+            print(_label_issues["label_score"])
+
+            # Display expected label issues
+            print("\nExpected label issues")
+            print(
+                _issues.assign(true_label=regression_data["true_y"][ids_of_union])
+                .query("true_label_error_idx")[
+                    [
+                        "is_label_issue",
+                        "label_score",
+                        "given_label",
+                        "predicted_label",
+                        "true_label",
+                    ]
+                ]
+                .sort_values("label_score")
+            )
+
+            for col in self.true_columns:
+                print(f"\nColumn: {col}")
+                print(_label_issues.query(col)["label_score"])
+
+            error_messages = [
+                (
+                    has_high_jaccard,
+                    f"- Jaccard similarity is too low. Should be > {expected_jaccard_similarity_bound}, got {float(intersection) / union}",
+                ),
+                (
+                    has_low_fpr,
+                    f"- FPR is too high. Should be less than or equal to {expected_fpr_bound}, got {fpr}",
+                ),
+            ]
+            error_msg = "The following test(s) failed for the default threshold:\n"
+            error_msg += "\n".join(
+                [msg for (test_passes, msg) in error_messages if not test_passes]
+            )
+
+            raise AssertionError(error_msg)
+
+    def test_regression_with_predictions_finds_label_issues(self, lab, regression_data):
+        """Test that the regression issue checks find 9 label issues, based on the
+        predictions of a model.
+
+        Instead of running a model, we use the ground-truth to emulate a perfect model's predictions.
+
+        Testing the default behavior, we expect to find some label issues with a given mean score.
+        Increasing a threshold for flagging issues will flag more issues, but won't change the score.
+        """
+
+        # Use ground-truth to emulate a perfect model's predictions
+        y_pred = regression_data["true_y"]
+        issue_types = {"label": {}}
+        lab.find_issues(pred_probs=y_pred, issue_types=issue_types)
+        summary = lab.get_issue_summary()
+
+        issues = lab.get_issues("label")
+
+        true_columns = [
+            "true_label_error_idx",
+            "true_outlier_idx",
+            "true_duplicate_idx",
+            "true_non_iid_idx",
+            "true_underperforming_group_idx",
+        ]
+        issue_ids = issues.query("is_label_issue").index
+        expected_issue_ids = regression_data["error_idx"]
+
+        # jaccard similarity
+        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
+        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
+        expected_jaccard_similarity_bound = 0.8
+        has_high_jaccard = float(intersection) / union > expected_jaccard_similarity_bound
+
+        # FPR
+        fpr = len(list(set(issue_ids).difference(set(expected_issue_ids)))) / len(issue_ids)
+        expected_fpr_bound = 0.0
+        has_no_fpr = fpr <= expected_fpr_bound
+
+        if not (has_high_jaccard and has_no_fpr):
+            # Work on showing more debugging information if end-to-end test fails
+            ids_of_union = list(set(issue_ids).union(set(expected_issue_ids)))
+            _issues = issues.loc[ids_of_union]
+            new_columns = [
+                "true_label_error_idx",
+                "true_outlier_idx",
+                "true_duplicate_idx",
+                "true_non_iid_idx",
+                "true_underperforming_group_idx",
+            ]
+            ids_keys = [
+                "error_idx",
+                "outlier_ids",
+                "duplicate_ids",
+                "non_iid_ids",
+                "underperforming_group_ids",
+            ]
+            for new_col, ids_key in zip(new_columns, ids_keys):
+                _issues[new_col] = _issues.index.isin(regression_data[ids_key])
+            # Display dataframe
+            print("Expected label issue ids", expected_issue_ids, "\n")
+            _label_issues = _issues.query("is_label_issue")
+            print(_label_issues)
+            for col in true_columns:
+                print(f"\nColumn: {col}")
+                print(_label_issues.query(col)["label_score"])
+
+            error_messages = [
+                (
+                    has_high_jaccard,
+                    f"- Jaccard similarity is too low. Should be > {expected_jaccard_similarity_bound}, got {float(intersection) / union}",
+                ),
+                (
+                    has_no_fpr,
+                    f"- FPR is too high. Should be less than or equal to {expected_fpr_bound}, got {fpr}",
+                ),
+            ]
+            error_msg = "The following test(s) failed for the default threshold:\n"
+            error_msg += "\n".join(
+                [msg for (test_passes, msg) in error_messages if not test_passes]
+            )
+            raise AssertionError(error_msg)
+
+        # Try running with a different threshold
+        lab.find_issues(pred_probs=y_pred, issue_types={"label": {"threshold": 0.2}})
+        issues = lab.get_issues("label")
+        issue_ids = issues.query("is_label_issue").index
+
+        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
+        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
+
+        different_threshold_jaccard = float(intersection) / union
+        assert float(intersection) / union > 0.3
+
+    def test_regression_with_model_and_features_finds_label_issues(self, lab, regression_data):
+        """Test that the regression issue checks find label issue with another model."""
+        from sklearn.ensemble import RandomForestRegressor, StackingRegressor
+        from sklearn.linear_model import RANSACRegressor
+        from sklearn.svm import LinearSVR
+
+        # Make an ensemble model of several models
+        estimators = [
+            ("ransac", RANSACRegressor(random_state=SEED)),
+            ("svr", LinearSVR(random_state=SEED)),
+        ]
+        model = StackingRegressor(
+            estimators=estimators, final_estimator=RandomForestRegressor(random_state=SEED)
+        )
+        X = regression_data["X"]
+        issue_types = {"label": {"clean_learning_kwargs": {"model": model, "seed": SEED}}}
+        lab.find_issues(features=X, issue_types=issue_types)
+
+        issues = lab.get_issues("label")
+        issue_ids = issues.query("is_label_issue").index
+        expected_issue_ids = regression_data[
+            "error_idx"
+        ]  # Set to 5% of the data, but random noise may be too small to detect
+
+        # jaccard similarity
+        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
+        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
+        expected_jaccard_similarity_bound = 0.3
+        has_high_jaccard = float(intersection) / union >= expected_jaccard_similarity_bound
+
+        # FPR
+        fpr = len(list(set(issue_ids).difference(set(expected_issue_ids)))) / len(issue_ids)
+        expected_fpr_bound = 0.3
+        has_low_fpr = fpr <= expected_fpr_bound
+
+        if not (has_high_jaccard and has_low_fpr):
+            ids_of_union = list(set(issue_ids).union(set(expected_issue_ids)))
+            _issues = issues.loc[ids_of_union]
+            ids_keys = [
+                "error_idx",
+                "outlier_ids",
+                "duplicate_ids",
+                "non_iid_ids",
+                "underperforming_group_ids",
+            ]
+            for new_col, ids_key in zip(self.true_columns, ids_keys):
+                _issues[new_col] = _issues.index.isin(regression_data[ids_key])
+            # Display dataframe
+            print("Expected label issue ids", expected_issue_ids, "\n")
+            _label_issues = _issues.query("is_label_issue")
+            print(_label_issues)
+
+            for col in self.true_columns:
+                print(f"\nColumn: {col}")
+                print(_label_issues.query(col)["label_score"])
+
+            error_messages = [
+                (
+                    has_high_jaccard,
+                    f"- Jaccard similarity is too low. Should be > {expected_jaccard_similarity_bound}, got {float(intersection) / union}",
+                ),
+                (
+                    has_low_fpr,
+                    f"- FPR is too high. Should be less than or equal to {expected_fpr_bound}, got {fpr}",
+                ),
+            ]
+            error_msg = "The following test(s) failed for the default threshold:\n"
+            error_msg += "\n".join(
+                [msg for (test_passes, msg) in error_messages if not test_passes]
+            )
+            raise AssertionError(error_msg)
+
+    @pytest.mark.parametrize(
+        "argument_name, data_key, expected_issue_types_in_summary",
+        [
+            ("features", "X", set(["label", "outlier", "near_duplicate", "non_iid", "null"])),
+            ("pred_probs", "true_y", set(["label"])),
+            ("knn_graph", "knn_graph", set(["outlier", "near_duplicate", "non_iid"])),
+        ],
+        ids=["features only", "pred_probs only", "knn_graph only"],
+    )
+    def test_find_issues_defaults(
+        self, lab, regression_data, argument_name, data_key, expected_issue_types_in_summary
+    ):
+        """Test that the regression issue checks find various issues with the default settings."""
+        input_dict = {argument_name: regression_data[data_key]}
+        lab.find_issues(**input_dict)
+        issue_summary = lab.get_issue_summary()
+        assert issue_summary["num_issues"].sum() > 0
+        assert set(issue_summary["issue_type"].values) == expected_issue_types_in_summary
+
+    @pytest.mark.parametrize(
+        "argument_name, data_key",
+        [
+            ("features", "X"),
+            ("knn_graph", "knn_graph"),
+        ],
+        ids=["features only", "knn_graph only"],
+    )
+    def test_find_outliers(self, lab, regression_data, argument_name, data_key):
+        """Test that the regression issue checks find 2 outlier issues."""
+
+        input_dict = {argument_name: regression_data[data_key]}
+        # Other tests are too sensitive to having more obvious outliers
+        issue_types = {
+            "outlier": {"threshold": 0.20}
+        }  # Lower threshold for a more conservative result
+        lab.find_issues(**input_dict, issue_types=issue_types)
+        lab.report()
+
+        issues = lab.get_issues("outlier")
+        issue_ids = issues.query("is_outlier_issue").index
+        expected_issue_ids = regression_data["outlier_ids"]
+
+        # jaccard similarity
+        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
+        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
+        assert float(intersection) / union >= 0.5
+
+        # FPR
+        fpr = len(list(set(issue_ids).difference(set(expected_issue_ids)))) / len(issue_ids)
+        assert fpr == pytest.approx(0.0, abs=0.01)
+
+    @pytest.mark.parametrize(
+        "argument_name, data_key",
+        [
+            ("features", "X"),
+            ("knn_graph", "knn_graph"),
+        ],
+        ids=["features only", "knn_graph only"],
+    )
+    def test_find_near_duplicates(self, lab, regression_data, argument_name, data_key):
+        """Test that the regression issue checks find 5 near-duplicate issues."""
+
+        input_dict = {argument_name: regression_data[data_key]}
+        issue_types = {"near_duplicate": {"threshold": 0.01}}
+        lab.find_issues(**input_dict, issue_types=issue_types)
+
+        issues = lab.get_issues("near_duplicate")
+        issue_ids = issues.query("is_near_duplicate_issue").index
+        expected_issue_ids = regression_data["duplicate_ids"]
+
+        # jaccard similarity
+        intersection = len(list(set(issue_ids).intersection(set(expected_issue_ids))))
+        union = len(set(issue_ids)) + len(set(expected_issue_ids)) - intersection
+        expected_jaccard_similarity_bound = 0.8
+        has_high_jaccard = float(intersection) / union > expected_jaccard_similarity_bound
+
+        # FPR
+        fpr = len(list(set(issue_ids).difference(set(expected_issue_ids)))) / len(issue_ids)
+        expected_fpr_bound = 0.2
+        has_low_fpr = fpr <= expected_fpr_bound
+
+        if not (has_high_jaccard and has_low_fpr):
+            ids_of_union = list(set(issue_ids).union(set(expected_issue_ids)))
+            _issues = issues.loc[ids_of_union]
+
+            ids_keys = [
+                "error_idx",
+                "outlier_ids",
+                "duplicate_ids",
+                "non_iid_ids",
+                "underperforming_group_ids",
+            ]
+            for new_col, ids_key in zip(self.true_columns, ids_keys):
+                _issues[new_col] = _issues.index.isin(regression_data[ids_key])
+            _near_duplicate_issues = _issues.query("is_near_duplicate_issue")
+            print(_near_duplicate_issues["near_duplicate_sets"])
+
+            for col in self.true_columns:
+                print(f"\nColumn: {col}")
+                print(_near_duplicate_issues.query(col)["near_duplicate_sets"])
+
+            error_messages = [
+                (
+                    has_high_jaccard,
+                    f"- Jaccard similarity is too low. Should be > {expected_jaccard_similarity_bound}, got {float(intersection) / union}",
+                ),
+                (
+                    has_low_fpr,
+                    f"- FPR is too high. Should be less than or equal to {expected_fpr_bound}, got {fpr}",
+                ),
+            ]
+            error_msg = "The following test(s) failed for the default threshold:\n"
+            error_msg += "\n".join(
+                [msg for (test_passes, msg) in error_messages if not test_passes]
+            )
+
+            raise AssertionError(error_msg)
+
+    @pytest.mark.parametrize(
+        "argument_name, data_key",
+        [
+            ("features", "X"),
+            ("knn_graph", "knn_graph"),
+        ],
+        ids=["features only", "knn_graph only"],
+    )
+    def test_find_non_iid(self, lab, regression_data, argument_name, data_key):
+        """Test that the regression issue checks find 30 non-iid issues."""
+
+        input_dict = {argument_name: regression_data[data_key]}
+        lab.find_issues(
+            **input_dict,
+            issue_types={
+                "non_iid": {"num_permutations": 1000, "seed": SEED, "significance_threshold": 0.3}
+            },
+        )
+
+        issues = lab.get_issues("non_iid")
+        issue_ids = issues.query("is_non_iid_issue").index
+        expected_issue_ids = regression_data["non_iid_ids"]
+
+        # omitting jaccard similarity and FPR for non-iid issues
+
+        # FPR
+        fp = list(set(issue_ids).difference(set(expected_issue_ids)))
+        fpr = len(fp) / len(issue_ids) if len(issue_ids) > 0 else 0.0
+        expected_fpr_bound = 0.0
+        has_low_fpr = fpr <= expected_fpr_bound
+
+        # AP@len(expected_issue_ids)
+        def average_precision_at_k(y_true, y_pred, k) -> float:
+            """Compute average precision at k."""
+            y_true = np.asarray(y_true)
+            y_pred = np.asarray(y_pred)
+            sort_indices = np.argsort(y_pred)[::-1]
+            y_true = y_true[sort_indices]
+            return float(np.mean(y_true[:k]))
+
+        ap_at_k = average_precision_at_k(
+            # y_true are the expected issues
+            y_true=np.isin(np.arange(len(issues)), expected_issue_ids),
+            # y_pred are the scores
+            y_pred=1 - issues["non_iid_score"],
+            k=len(expected_issue_ids),
+        )
+        expected_ap_at_k_bound = 0.8
+        has_high_enough_ap_at_k = ap_at_k >= expected_ap_at_k_bound
+
+        if not (has_low_fpr and has_high_enough_ap_at_k):
+            ids_of_union = list(set(issue_ids).union(set(expected_issue_ids)))
+            _issues = issues.loc[ids_of_union]
+            ids_keys = [
+                "error_idx",
+                "outlier_ids",
+                "duplicate_ids",
+                "non_iid_ids",
+                "underperforming_group_ids",
+            ]
+            for new_col, ids_key in zip(self.true_columns, ids_keys):
+                _issues[new_col] = _issues.index.isin(regression_data[ids_key])
+            _non_iid_issues = _issues.query("is_non_iid_issue")
+            print(
+                _issues[["is_non_iid_issue", "non_iid_score"]]
+                .sort_values("non_iid_score")
+                .head(len(expected_issue_ids))
+            )
+
+            for col in self.true_columns:
+                print(f"\nColumn: {col}")
+                # print(_non_iid_issues.query(col)["non_iid_sets"])
+
+            error_messages = [
+                (
+                    has_low_fpr,
+                    f"- FPR is too high. Should be less than or equal to {expected_fpr_bound}, got {fpr}",
+                ),
+                (
+                    has_high_enough_ap_at_k,
+                    f"- AP@{len(expected_issue_ids)} should be greater than or equal to {expected_ap_at_k_bound}, got {ap_at_k}",
+                ),
+            ]
+            error_msg = "The following test(s) failed for the default threshold:\n"
+            error_msg += "\n".join(
+                [msg for (test_passes, msg) in error_messages if not test_passes]
+            )
+
+            raise AssertionError(error_msg)
+
+    def test_find_null_issues(self, lab, regression_data):
+        """Test that the regression issue checks find 0 null issues."""
+        X = regression_data["X"]
+        lab.find_issues(features=X, issue_types={"null": {}})
+        summary = lab.get_issue_summary("null")
+        assert summary["num_issues"].values[0] == 0
+
+        rand_ids = np.random.choice(X.shape[0], 10, replace=False)
+        for i in rand_ids:
+            j = np.random.choice(X.shape[1], 1, replace=False)
+            X[i, j] = np.nan
+
+        rand_ids_full = np.random.choice(X.shape[0], 3, replace=False)
+        for i in rand_ids_full:
+            X[i, :] = np.nan
+        lab.find_issues(features=X, issue_types={"null": {}})
+        issues = lab.get_issues("null")
+        null_issues = issues.query("is_null_issue")
+        assert set(rand_ids_full) == set(null_issues.index.tolist())

tests/datalab/issue_manager/conftest.py

@@ -0,0 +1,89 @@
+from hypothesis import strategies as st
+import numpy as np
+from scipy.sparse import csr_matrix
+
+
+@st.composite
+def knn_graph_strategy(draw, num_samples, k_neighbors, min_distance=0.0, max_distance=100.0):
+    """
+    Generate a K-nearest neighbors (KNN) graph based on the given parameters.
+
+    Parameters
+    ----------
+    draw: A function used to draw values from search strategies.
+
+    num_samples (int or SearchStrategy): The number of samples in the graph.
+        If a SearchStrategy is provided, a value will be drawn from it.
+
+    k_neighbors (int or SearchStrategy): The number of nearest neighbors to consider for each sample.
+        If a SearchStrategy is provided, a value will be drawn from it.
+
+    Returns
+    -------
+    knn_graph : csr_matrix
+        The KNN graph represented as a sparse matrix.
+
+    Notes
+    -----
+    - The KNN graph is generated based on a symmetric distance matrix.
+    - The distance matrix is computed using randomly generated upper triangle values.
+    - The diagonal of the distance matrix is set to infinity to avoid selecting a point as its own neighbor.
+    - The K-nearest neighbors are computed based on the distance matrix.
+    - The resulting KNN graph is returned as a sparse matrix in csr format.
+    - The number of samples must be greater than the number of neighbors.
+    - The KNN graph is not guaranteed to be connected (i.e. there may be isolated subgraphs).
+    - The KNN graph is a directed graph (i.e. the edges are not symmetric).
+    - The neighbors are sorted by distance in the CSR-formatted sparse matrix,
+        so the first neighbor is the closest neighbor.
+    """
+    # If the argument is a strategy, draw a value from it.
+    if isinstance(num_samples, st.SearchStrategy):
+        num_samples = draw(num_samples)
+
+    if isinstance(k_neighbors, st.SearchStrategy):
+        k_neighbors = draw(k_neighbors)
+
+    # Generate a symmetric distance matrix
+    upper_triangle = [
+        draw(
+            st.lists(
+                st.floats(
+                    min_value=min_distance,
+                    max_value=max_distance,
+                    allow_nan=False,
+                    allow_infinity=False,
+                    allow_subnormal=False,
+                ),
+                min_size=i,
+                max_size=i,
+                unique=True,
+            )
+        )
+        for i in range(1, num_samples + 1)
+    ]
+
+    distance_matrix = np.zeros((num_samples, num_samples))
+    for i, row in enumerate(upper_triangle):
+        distance_matrix[i, : i + 1] = row
+        distance_matrix[: i + 1, i] = row
+
+    np.fill_diagonal(
+        distance_matrix, np.inf
+    )  # To ensure we don't select a point as its own neighbor
+
+    # Compute k-nearest neighbors based on the distance matrix
+    sorted_indices = np.argsort(distance_matrix, axis=1)
+    kneighbor_indices = sorted_indices[:, :k_neighbors]
+    kneighbor_distances = np.array(
+        [distance_matrix[i, kneighbor_indices[i]] for i in range(num_samples)]
+    )
+
+    knn_graph = csr_matrix(
+        (
+            kneighbor_distances.flatten(),
+            kneighbor_indices.flatten(),
+            np.arange(0, (kneighbor_distances.shape[0] * k_neighbors + 1), k_neighbors),
+        ),
+        shape=(num_samples, num_samples),
+    )
+    return knn_graph

tests/datalab/issue_manager/multilabel/test_label.py

@@ -0,0 +1,48 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from cleanlab import Datalab
+from cleanlab.datalab.internal.issue_manager.multilabel.label import MultilabelIssueManager
+from cleanlab.internal.multilabel_utils import onehot2int
+
+
+class TestLabelIssueManager:
+    @pytest.fixture
+    def data(self):
+        # True labels of multilabel dataset
+        np.random.seed(10)
+        true_y = np.array(
+            [[0, 1, 0], [0, 1, 1], [1, 1, 1], [1, 0, 1], [0, 1, 1], [1, 1, 0], [1, 1, 0]]
+        )
+        pred_probs = np.full_like(true_y, fill_value=0.1, dtype=float)
+        pred_probs[true_y == 1] = 0.9
+        # Flip labels of some rows to add noise
+        candidate_rows = np.where(true_y.sum(axis=1) < 3)[0]
+        noisy_rows = np.random.choice(candidate_rows, 2, replace=False)
+        noisy_y = true_y.copy()
+        noisy_y[noisy_rows] = 1 - true_y[noisy_rows]
+        labels = onehot2int(noisy_y)
+        return {"labels": labels, "pred_probs": pred_probs}
+
+    @pytest.fixture
+    def issue_manager(self, data):
+        labels = data["labels"]
+        lab = Datalab({"labels": labels}, task="multilabel", label_name="labels")
+        return MultilabelIssueManager(datalab=lab)
+
+    def test_find_issues(self, data, issue_manager):
+        """Test that the find_issues method works."""
+        pred_probs = data["pred_probs"]
+        issue_manager.find_issues(pred_probs=pred_probs)
+        issues, summary, info = issue_manager.issues, issue_manager.summary, issue_manager.info
+        assert isinstance(issues, pd.DataFrame), "Issues should be a dataframe"
+        assert isinstance(summary, pd.DataFrame), "Summary should be a dataframe"
+        assert summary["issue_type"].values[0] == "label"
+        assert issues.index[issues["is_label_issue"]].tolist() == [3, 6]
+        assert pytest.approx(summary["score"].values[0], abs=1e-3) == 0.6714
+        assert isinstance(info, dict), "Info should be a dict"
+
+        issue_manager.find_issues(pred_probs=pred_probs, frac_noise=0.5)
+        issues = issue_manager.issues
+        assert issues.index[issues["is_label_issue"]].tolist() == [3]

tests/datalab/issue_manager/regression/test_label.py

@@ -0,0 +1,145 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from cleanlab import Datalab
+from cleanlab.datalab.internal.issue_manager.regression.label import RegressionLabelIssueManager
+from cleanlab.datalab.internal.task import Task
+
+
+def ground_truth_target_function(x):
+    return 10 * x + 1
+
+
+class TestRegressionLabelIssueManager:
+    def test_manager_found_in_registry(self):
+        from cleanlab.datalab.internal.issue_manager_factory import REGISTRY
+
+        error_msg = (
+            "RegressionLabelIssueManager should be registered to the regression task as 'label'"
+        )
+        assert REGISTRY[Task.REGRESSION].get("label") == RegressionLabelIssueManager, error_msg
+
+    @pytest.fixture
+    def features(self):
+        # 1 feature, 7 points
+        return np.array([0.1, 0.2, 0.3, 0.35, 0.4, 0.45, 0.5]).reshape(-1, 1)
+
+    @pytest.fixture
+    def regression_lab(self, features):
+        y = ground_truth_target_function(features)
+        # Flip the sign of the point x=0.4
+        y[features == 0.4] *= -1
+        y = y.ravel()
+        return Datalab({"y": y}, label_name="y", task="regression")
+
+    @pytest.fixture
+    def issue_manager(self, regression_lab):
+        return RegressionLabelIssueManager(datalab=regression_lab)
+
+    def test_find_issues_with_features(self, issue_manager, features):
+        issue_manager.find_issues(features=features)
+        issues = issue_manager.issues
+        assert isinstance(issues, pd.DataFrame), "Issues should be a dataframe"
+        expected_issue_mask = features.ravel() == 0.4
+        assert sum(expected_issue_mask) == 1, "There should be exactly one issue"
+
+        np.testing.assert_array_equal(issues["is_label_issue"].values, expected_issue_mask)
+        # Assert that he minimum score "label_score" is at the correct index
+        index_of_error = np.where(expected_issue_mask)[0][0]
+        assert issues["label_score"].values.argmin() == index_of_error
+
+    def test_init_with_model(self, issue_manager):
+        from sklearn.neighbors import KNeighborsRegressor
+
+        model = KNeighborsRegressor(n_neighbors=2)
+        assert issue_manager.cl.model != model
+
+        # Passing in a model to the constructor should set the cl.model field
+        clean_learning_kwargs = {"model": model}
+        lab = issue_manager.datalab
+        new_issue_manager = RegressionLabelIssueManager(
+            datalab=lab, clean_learning_kwargs=clean_learning_kwargs
+        )
+        assert new_issue_manager.cl.model == model
+
+    @pytest.fixture
+    def predictions(self, features):
+        y_ground_truth = ground_truth_target_function(features).ravel()
+        noise = 0.1 * np.random.randn(len(y_ground_truth))
+        return y_ground_truth + noise
+
+    def test_raises_find_issues_error_without_valid_inputs(self, issue_manager):
+        with pytest.raises(ValueError) as e:
+            expected_error_msg = (
+                "Regression requires numerical `features` or `predictions` "
+                "to be passed in as an argument to `find_issues`."
+            )
+            issue_manager.find_issues()
+            assert expected_error_msg in str(e)
+
+    def test_find_issue_with_predictions(self, issue_manager, features, predictions):
+        issue_manager.find_issues(predictions=predictions)
+        issues = issue_manager.issues
+        assert isinstance(issues, pd.DataFrame), "Issues should be a dataframe"
+        expected_issue_mask = features.ravel() == 0.4
+        assert sum(expected_issue_mask) == 1, "There should be exactly one issue"
+
+        np.testing.assert_array_equal(issues["is_label_issue"].values, expected_issue_mask)
+        # Assert that he minimum score "label_score" is at the correct index
+        index_of_error = np.where(expected_issue_mask)[0][0]
+        assert issues["label_score"].values.argmin() == index_of_error
+
+
+class TestRegressionLabelIssueManagerIntegration:
+    """This class contains tests for the find_issues method with a CleanLearning
+    object that behaves deterministically. This is useful to run a "regression"-test on
+    the results computed by the find_issues method.
+    The test dataset is a random toy regression dataset with 5 features and 100 samples.
+    The ground truth is a linear function of the first feature plus a bias defined in the
+    class attribute BIAS.
+    The ground truth is used to emulate a perfect model and compute the expected score
+    for the label issue detection. The gaussian noise contributes to lower label quality
+    scores.
+    """
+
+    BIAS = 1.0
+
+    @pytest.fixture()
+    def regression_dataset(self):
+        """For integration tests, a simple regression dataset is simpler than
+        a tiny, hand-crafted one."""
+        from sklearn.datasets import make_regression
+
+        # Return coefficients as well for testing purposes,
+        # interpret as ground truth
+        X, y, coef = make_regression(
+            n_samples=100,
+            n_features=5,
+            n_informative=1,
+            n_targets=1,
+            bias=self.BIAS,
+            noise=0.1,
+            random_state=0,
+            coef=True,
+        )
+        return X, y, coef
+
+    @pytest.fixture()
+    def issue_manager(self, regression_dataset):
+        _, y, _ = regression_dataset
+        lab = Datalab({"y": y}, label_name="y", task="regression")
+        return RegressionLabelIssueManager(datalab=lab, clean_learning_kwargs={"seed": 0})
+
+    def test_find_issues_with_features(self, regression_dataset, issue_manager):
+        X, _, _ = regression_dataset
+        issue_manager.find_issues(features=X)
+        summary = issue_manager.summary
+        assert np.isclose(summary["score"], 0.425874, atol=1e-5)
+
+    def test_find_issues_with_predictions(self, regression_dataset, issue_manager):
+        X, _, coef = regression_dataset
+        y_pred = X @ coef + self.BIAS
+        issue_manager.find_issues(predictions=y_pred)
+        summary = issue_manager.summary
+        assert np.isclose(summary["score"], 0.361287, atol=1e-5)

tests/datalab/issue_manager/test_data_valuation.py

@@ -0,0 +1,85 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from cleanlab.datalab.internal.issue_manager.outlier import OutlierIssueManager
+from cleanlab.datalab.internal.issue_manager.data_valuation import DataValuationIssueManager
+from cleanlab.internal.neighbor.knn_graph import create_knn_graph_and_index
+
+SEED = 42
+
+
+class TestDataValuationIssueManager:
+    @pytest.fixture
+    def issue_manager(self, lab):
+        return DataValuationIssueManager(datalab=lab, k=3)
+
+    @pytest.fixture
+    def outlier_issue_manager(self, lab):
+        return OutlierIssueManager(datalab=lab, k=3)
+
+    @pytest.fixture
+    def embeddings(self, lab):
+        np.random.seed(SEED)
+        embeddings_array = 0.5 + 0.1 * np.random.rand(lab.get_info("statistics")["num_examples"], 2)
+        embeddings_array[4, :] = -1
+        return {"embedding": embeddings_array}
+
+    def test_find_issues_with_input(self, issue_manager, embeddings):
+        knn_graph, _ = create_knn_graph_and_index(
+            embeddings["embedding"],
+            n_neighbors=3,
+        )
+        issue_manager.find_issues(knn_graph=knn_graph)
+        issues, summary, info = issue_manager.issues, issue_manager.summary, issue_manager.info
+        assert isinstance(issues, pd.DataFrame), "Issues should be a dataframe"
+
+        assert isinstance(summary, pd.DataFrame), "Summary should be a dataframe"
+        assert summary["issue_type"].values[0] == "data_valuation"
+
+        assert isinstance(info, dict), "Info should be a dict"
+        assert isinstance(issues, pd.DataFrame), "Issues should be a dataframe"
+        info_keys = info.keys()
+        expected_keys = [
+            "num_low_valuation_issues",
+            "average_data_valuation",
+        ]
+        assert all(
+            [key in info_keys for key in expected_keys]
+        ), f"Info should have the right keys, but is missing {set(expected_keys) - set(info_keys)}"
+
+    def test_find_issues_with_stats(self, issue_manager, embeddings):
+        issue_manager.datalab.find_issues(
+            features=embeddings["embedding"], issue_types={"outlier": {"k": 3}}
+        )
+        issue_manager.find_issues(issue_types={"data_valuation": {"k": 3}})
+        issues, summary, info = issue_manager.issues, issue_manager.summary, issue_manager.info
+        assert isinstance(issues, pd.DataFrame), "Issues should be a dataframe"
+
+        assert isinstance(summary, pd.DataFrame), "Summary should be a dataframe"
+        assert summary["issue_type"].values[0] == "data_valuation"
+
+        assert isinstance(info, dict), "Info should be a dict"
+        assert isinstance(issues, pd.DataFrame), "Issues should be a dataframe"
+        info_keys = info.keys()
+        expected_keys = [
+            "num_low_valuation_issues",
+            "average_data_valuation",
+        ]
+        assert all(
+            [key in info_keys for key in expected_keys]
+        ), f"Info should have the right keys, but is missing {set(expected_keys) - set(info_keys)}"
+
+    def test_get_larger_k_than_knn_graph(self, issue_manager, embeddings, outlier_issue_manager):
+        outlier_issue_manager.find_issues(features=embeddings["embedding"])
+        knn_graph, _ = create_knn_graph_and_index(
+            embeddings["embedding"],
+            n_neighbors=3,
+        )
+        issue_manager.k = 4
+        expected_error_msg = (
+            "The provided knn graph has 3 neighbors, which is less than the required 4 neighbors. "
+            "Please ensure that the knn graph you provide has at least as many neighbors as the required value of k."
+        )
+        with pytest.raises(ValueError, match=expected_error_msg):
+            issue_manager.find_issues(knn_graph=knn_graph)

tests/datalab/issue_manager/test_duplicate.py

@@ -0,0 +1,258 @@
+import numpy as np
+import pytest
+from hypothesis import HealthCheck, assume, given, settings, strategies as st
+from hypothesis.strategies import composite
+from hypothesis.extra.numpy import arrays
+
+
+from cleanlab import Datalab
+from cleanlab.datalab.internal.issue_manager.duplicate import NearDuplicateIssueManager
+
+from .conftest import knn_graph_strategy
+
+SEED = 42
+
+
+@composite
+def embeddings_strategy(draw):
+    shape_strategy = st.tuples(
+        st.integers(min_value=3, max_value=20), st.integers(min_value=2, max_value=2)
+    )
+    element_strategy = st.floats(
+        min_value=0.0, max_value=1.0, allow_nan=False, allow_infinity=False
+    )
+    embeddings = draw(
+        arrays(
+            dtype=np.float64,
+            shape=shape_strategy,
+            elements=element_strategy,
+            unique=True,
+        )
+    )
+    return embeddings
+
+
+class TestNearDuplicateIssueManager:
+    @pytest.fixture
+    def embeddings(self, lab):
+        np.random.seed(SEED)
+        embeddings_array = 0.5 + 0.1 * np.random.rand(lab.get_info("statistics")["num_examples"], 2)
+        embeddings_array[4, :] = (
+            embeddings_array[3, :] + np.random.rand(embeddings_array.shape[1]) * 0.001
+        )
+        return {"embedding": embeddings_array}
+
+    @pytest.fixture
+    def issue_manager(self, lab, embeddings, monkeypatch):
+        mock_data = lab.data.from_dict({**lab.data.to_dict(), **embeddings})
+        monkeypatch.setattr(lab, "data", mock_data)
+        return NearDuplicateIssueManager(
+            datalab=lab,
+            metric="euclidean",
+            k=2,
+        )
+
+    def test_init(self, lab, issue_manager):
+        assert issue_manager.datalab == lab
+        assert issue_manager.metric == "euclidean"
+        assert issue_manager.k == 2
+        assert issue_manager.threshold == 0.13
+
+        issue_manager = NearDuplicateIssueManager(
+            datalab=lab,
+            threshold=0.1,
+        )
+        assert issue_manager.threshold == 0.1
+
+    def test_find_issues(self, issue_manager, embeddings):
+        issue_manager.find_issues(features=embeddings["embedding"])
+        issues, summary, info = issue_manager.issues, issue_manager.summary, issue_manager.info
+        expected_issue_mask = np.array([False] * 3 + [True] * 2)
+        assert np.all(
+            issues["is_near_duplicate_issue"] == expected_issue_mask
+        ), "Issue mask should be correct"
+        assert summary["issue_type"][0] == "near_duplicate"
+        assert summary["score"][0] == pytest.approx(expected=0.4734458, abs=1e-7)
+
+        assert (
+            info.get("near_duplicate_sets", None) is not None
+        ), "Should have sets of near duplicates"
+
+        new_issue_manager = NearDuplicateIssueManager(
+            datalab=issue_manager.datalab,
+            metric="euclidean",
+            k=2,
+            threshold=0.1,
+        )
+        new_issue_manager.find_issues(features=embeddings["embedding"])
+
+    def test_scores_of_examples_with_issues_are_smaller_than_those_without(
+        self, issue_manager, embeddings
+    ):
+        # TODO: Turn this into a property-based test
+        issue_manager.find_issues(features=embeddings["embedding"])
+        is_issue = issue_manager.issues["is_near_duplicate_issue"]
+        scores = issue_manager.issues["near_duplicate_score"]
+        max_issue_score = np.max(scores[is_issue])
+        min_non_issue_score = np.min(scores[~is_issue])
+        assert max_issue_score < min_non_issue_score
+
+    def test_report(self, issue_manager, embeddings):
+        issue_manager.find_issues(features=embeddings["embedding"])
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+        )
+        assert isinstance(report, str)
+        assert (
+            "------------------ near_duplicate issues -------------------\n\n"
+            "Number of examples with this issue:"
+        ) in report
+
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+            verbosity=3,
+        )
+        assert "Additional Information: " in report
+
+    @given(embeddings=embeddings_strategy())
+    @settings(deadline=800)
+    def test_near_duplicate_sets(self, embeddings):
+        data = {"metadata": ["" for _ in range(len(embeddings))]}
+        lab = Datalab(data)
+        issue_manager = NearDuplicateIssueManager(
+            datalab=lab,
+            metric="euclidean",
+            k=2,
+        )
+        embeddings = np.array(embeddings)
+        issue_manager.find_issues(features=embeddings)
+        near_duplicate_sets = issue_manager.info["near_duplicate_sets"]
+        issues = issue_manager.issues["is_near_duplicate_issue"]
+
+        # Test: Near duplicates are symmetric
+        all_symmetric = all(
+            i in near_duplicate_sets[j]
+            for i, near_duplicates in enumerate(near_duplicate_sets)
+            for j in near_duplicates
+        )
+        assert all_symmetric, "Some near duplicate sets are not symmetric"
+
+        # Test: Near duplicate sets for issues
+        all_non_issues_have_empty_near_duplicate_sets = all(
+            len(near_duplicate_set) == 0
+            for i, near_duplicate_set in enumerate(near_duplicate_sets)
+            if not issues[i]
+        )
+        assert (
+            all_non_issues_have_empty_near_duplicate_sets
+        ), "Non-issue examples should not have near duplicate sets"
+        all_issues_have_non_empty_near_duplicate_sets = all(
+            len(near_duplicate_set) > 0
+            for i, near_duplicate_set in enumerate(near_duplicate_sets)
+            if issues[i]
+        )
+        assert (
+            all_issues_have_non_empty_near_duplicate_sets
+        ), "Issue examples should have near duplicate sets"
+
+
+def build_issue_manager(
+    draw, num_samples_strategy, k_neighbors_strategy, with_issues=False, threshold=None
+):
+    """Create a random knn_graph with the given number of samples and k neighbors.
+    Run the NearDuplicateIssueManager on the knn_graph and return the issue manager.
+    A threshold can be provided to control the number of issues for small graphs.
+    A with_issues flag can be provided to control whether the issue manager should have issues.
+    """
+
+    if with_issues:
+        knn_graph = draw(
+            knn_graph_strategy(num_samples=num_samples_strategy, k_neighbors=k_neighbors_strategy)
+        )
+    else:
+        knn_graph = draw(
+            knn_graph_strategy(
+                num_samples=num_samples_strategy, k_neighbors=k_neighbors_strategy, min_distance=0.1
+            )
+        )
+
+    lab = Datalab(data={})
+    inputs = {"datalab": lab, "threshold": threshold}
+    inputs = {k: v for k, v in inputs.items() if v is not None}
+    issue_manager = NearDuplicateIssueManager(**inputs)
+    issue_manager.find_issues(knn_graph=knn_graph)
+    issues = issue_manager.issues["is_near_duplicate_issue"]
+
+    if with_issues:
+        assume(any(issues))
+    else:
+        assume(not any(issues))
+    return issue_manager
+
+
+@st.composite
+def no_issue_issue_manager_strategy(draw):
+    """Strategy for generating NearDuplicateIssueManagers with no issues."""
+    return build_issue_manager(
+        draw,
+        st.integers(min_value=10, max_value=50),
+        st.integers(min_value=2, max_value=5),
+        with_issues=False,
+        threshold=0.0001,
+    )
+
+
+@st.composite
+def issue_manager_with_issues_strategy(draw):
+    """Strategy for generating NearDuplicateIssueManagers with issues."""
+    return build_issue_manager(
+        draw,
+        st.integers(min_value=10, max_value=20),
+        st.integers(min_value=2, max_value=5),
+        with_issues=True,
+        threshold=0.9,
+    )
+
+
+class TestNearDuplicateSets:
+    """Property-based tests properties of near duplicate sets found in a knn graph."""
+
+    @pytest.mark.slow
+    @given(issue_manager=no_issue_issue_manager_strategy())
+    @settings(
+        deadline=800, suppress_health_check=[HealthCheck.too_slow, HealthCheck.data_too_large]
+    )
+    def test_near_duplicate_sets_empty_if_no_issue_next(self, issue_manager):
+        near_duplicate_sets = issue_manager.info["near_duplicate_sets"]
+        assert all(len(near_duplicate_set) == 0 for near_duplicate_set in near_duplicate_sets)
+
+    @given(issue_manager=issue_manager_with_issues_strategy())
+    @settings(deadline=800, max_examples=1000, suppress_health_check=[HealthCheck.too_slow])
+    def test_symmetric_and_flagged_consistency(self, issue_manager):
+        near_duplicate_sets = issue_manager.info["near_duplicate_sets"]
+        issues = issue_manager.issues["is_near_duplicate_issue"]
+
+        # Test symmetry: If A is in near_duplicate_set of B, then B should be in near_duplicate_set of A.
+        for i, near_duplicates in enumerate(near_duplicate_sets):
+            for j in near_duplicates:
+                assert (
+                    i in near_duplicate_sets[j]
+                ), f"Example {j} is in near_duplicate_set of {i}, but not vice versa"
+
+        # Test consistency of flags with near_duplicate_sets
+        for i, near_duplicate_set in enumerate(near_duplicate_sets):
+            if issues[i]:
+                # Near duplicate sets of flagged examples should not be empty
+                assert (
+                    len(near_duplicate_set) > 0
+                ), f"Near duplicate set of flagged example {i} is empty"
+
+                # Check if all examples in the near_duplicate_set of a flagged example are also flagged
+                flagged_in_set = [issues[j] for j in near_duplicate_set]
+                assert all(
+                    flagged_in_set
+                ), f"Example {i} is flagged as near_duplicate but some examples in its near_duplicate_set are not flagged"

tests/datalab/issue_manager/test_identifier_column.py

@@ -0,0 +1,138 @@
+import numpy as np
+import pandas as pd
+import pytest
+from cleanlab.datalab.internal.issue_manager.identifier_column import IdentifierColumnIssueManager
+
+
+class TestIdentifierColumnIssueManager:
+    @pytest.fixture
+    def issue_manager(self, lab):
+        return IdentifierColumnIssueManager(datalab=lab)
+
+    @pytest.mark.parametrize(
+        "arr, expected_output",
+        [
+            (np.array([1, 2, 3, 4, 5]), True),
+            (np.array([1, 1, 2, 2, 3, 3, 5]), False),
+            (np.array([1, 1, 3, 4, 5, 8, 10]), False),
+            (np.array([0, 0, 0, 0, 0, 0, 0]), False),
+            (np.array([4, 5, 5, 6, 7, 8, 9, 10]), True),
+            (np.array([1, 3, 4, 4, 5, 6, 7, -1]), False),
+            (np.array([2, 1, 3, 5, 6, 4]), True),
+            (np.array([-1, -3, -2, -4, 0]), True),
+            (np.array([]), False),
+            (np.array([0, 0, 0]), False),
+        ],
+    )
+    def test_is_sequential(self, issue_manager, arr, expected_output):
+        assert issue_manager._is_sequential(arr) == expected_output
+
+    @pytest.mark.parametrize(
+        "features, expected_prepared_features",
+        [
+            (np.array([[1, 2, 3], [4, 5, 6]]), np.array([[1, 4], [2, 5], [3, 6]])),
+            (
+                pd.DataFrame({"A": [1, 4], "B": [2, 5], "C": [3, 6]}),
+                [np.array([1, 4]), np.array([2, 5]), np.array([3, 6])],
+            ),
+            (
+                pd.DataFrame({"A": [1, 4], "B": [2.0, 5.0], "C": [3, 6]}),
+                [np.array([1, 4]), np.array([2.0, 5.0]), np.array([3, 6])],
+            ),
+            (
+                pd.DataFrame({"A": [1, 4], "B": [2, 5], "C": ["3", "6"]}),
+                [np.array([1, 4]), np.array([2, 5]), np.array(["3", "6"], dtype=str)],
+            ),
+            ([[1, 4], [2, 5], [3, 6]], [np.array([1, 4]), np.array([2, 5]), np.array([3, 6])]),
+            (
+                {"A": [1, 4], "B": [2, 5], "C": [3, 6]},
+                [np.array([1, 4]), np.array([2, 5]), np.array([3, 6])],
+            ),
+            (
+                {"A": [1, 4], "B": [2.0, 5.0], "C": [3, 6], "D": ["a", "b"]},
+                [
+                    np.array([1, 4]),
+                    np.array([2.0, 5.0]),
+                    np.array([3, 6]),
+                    np.array(["a", "b"], dtype=str),
+                ],
+            ),
+        ],
+    )
+    def test_prepare_features(self, issue_manager, features, expected_prepared_features):
+        prepared_features = issue_manager._prepare_features(features)
+        assert np.array_equal(prepared_features, expected_prepared_features)
+
+    @pytest.mark.parametrize(
+        "features, expected_indices, expected_is_identifier_column",
+        [
+            (np.array([[1, 2, 3], [4, 5, 2]]), [2], 0.0),
+            (np.array([[1, 2, 3], [1, 3, 2]]), [1, 2], 0.0),
+            (
+                np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
+                [],
+                1.0,
+            ),
+            (
+                np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]]),
+                [],
+                1.0,
+            ),
+            (
+                np.array([[0, 2, 3], [-1, 5, 6], [-2, 2, 3]]),
+                [0],
+                0.0,
+            ),
+            (
+                np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]),
+                [],
+                1.0,
+            ),
+            (np.array([[1, 2, 7], [4, 3, 8], [7, 4, 9], [10, 5, 10]]), [1, 2], 0.0),
+        ],
+    )
+    def test_find_issues(
+        self, issue_manager, features, expected_indices, expected_is_identifier_column
+    ):
+        issue_manager.find_issues(features)
+        print(f"summary: {issue_manager.summary['score'].values[0]}")
+        # print type of score
+        score = issue_manager.summary["score"].values[0]
+        print(f"score type: {type(score)}")
+        print(f"num_identifier_columns: {issue_manager.info['num_identifier_columns']}")
+        assert issue_manager.summary["score"].values[0] == expected_is_identifier_column
+        assert issue_manager.info["num_identifier_columns"] == len(expected_indices)
+        assert np.array_equal(issue_manager.info["identifier_columns"], expected_indices)
+
+    @pytest.mark.parametrize(
+        "features, expected_is_identifier_column_issue, expected_is_identifier_column",
+        [
+            (np.array([[1, 2, 3], [4, 5, 2]]), np.array([False, False]), [1.0, 1.0]),
+            (np.array([[1, 2, 3], [1, 3, 2]]), np.array([False, False]), [1.0, 1.0]),
+            (
+                np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
+                np.array([False, False, False]),
+                [1.0, 1.0, 1.0],
+            ),
+            (
+                np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]]),
+                np.array([False, False, False]),
+                [1.0, 1.0, 1.0],
+            ),
+        ],
+    )
+    def test_issue_attribute(
+        self,
+        issue_manager,
+        features,
+        expected_is_identifier_column_issue,
+        expected_is_identifier_column,
+    ):
+        issue_manager.find_issues(features)
+        assert np.array_equal(
+            issue_manager.issues[f"is_{issue_manager.issue_name}_issue"],
+            expected_is_identifier_column_issue,
+        )
+        assert np.array_equal(
+            issue_manager.issues[issue_manager.issue_score_key], expected_is_identifier_column
+        )

tests/datalab/issue_manager/test_imbalance.py

@@ -0,0 +1,106 @@
+import numpy as np
+import pytest
+
+from cleanlab.datalab.internal.issue_manager.imbalance import ClassImbalanceIssueManager
+
+SEED = 42
+
+
+class TestClassImbalanceIssueManager:
+    @pytest.fixture
+    def labels(self, lab):
+        np.random.seed(SEED)
+        K = lab.get_info("statistics")["num_classes"]
+        N = lab.get_info("statistics")["num_examples"] * 20
+        labels = np.random.choice(np.arange(K - 1), size=N, p=[0.5] * (K - 1))
+        labels[0] = K - 1  # Rare class
+        return labels
+
+    @pytest.fixture
+    def create_issue_manager(self, lab, labels, monkeypatch):
+        def manager(labels=labels):
+            monkeypatch.setattr(lab._labels, "labels", labels)
+            return ClassImbalanceIssueManager(datalab=lab, threshold=0.1)
+
+        return manager
+
+    def test_find_issues(self, create_issue_manager, labels):
+        N = len(labels)
+        issue_manager = create_issue_manager()
+        issue_manager.find_issues()
+        issues, summary = issue_manager.issues, issue_manager.summary
+        assert np.sum(issues["is_class_imbalance_issue"]) == 1
+        expected_issue_mask = np.array([True] + [False] * (N - 1))
+        assert np.all(
+            issues["is_class_imbalance_issue"] == expected_issue_mask
+        ), "Issue mask should be correct"
+        expected_scores = np.array([0.01] + [1.0] * (N - 1))
+        np.testing.assert_allclose(
+            issues["class_imbalance_score"], expected_scores, err_msg="Scores should be correct"
+        )
+        assert summary["issue_type"][0] == "class_imbalance"
+        assert summary["score"][0] == pytest.approx(0.01, abs=0.01)
+
+    def test_find_issues_no_imbalance(self, labels, create_issue_manager):
+        N = len(labels)
+        labels[0] = 0
+        issue_manager = create_issue_manager(labels)
+        issue_manager.find_issues()
+        issues, summary = issue_manager.issues, issue_manager.summary
+        assert np.sum(issues["is_class_imbalance_issue"]) == 0
+        assert np.all(
+            issues["is_class_imbalance_issue"] == np.full(N, False)
+        ), "Issue mask should be correct"
+        scores = issues["class_imbalance_score"]
+        expected_scores = np.ones_like(scores)
+        expected_scores[labels == 1] = 0.47  # Rare class proportion
+        np.testing.assert_allclose(scores, expected_scores, err_msg="Scores should be correct")
+        assert summary["issue_type"][0] == "class_imbalance"
+        assert summary["score"][0] == pytest.approx(0.47, abs=0.01)
+
+    def test_find_issues_more_imbalance(self, lab, labels, create_issue_manager):
+        K = lab.get_info("statistics")["num_classes"]
+        N = len(labels)
+        labels[labels == K - 2] = 0
+        labels[1:3] = K - 2
+        issue_manager = create_issue_manager(labels)
+        issue_manager.find_issues()
+        issues, summary = issue_manager.issues, issue_manager.summary
+        assert np.sum(issues["is_class_imbalance_issue"]) == 1
+        expected_issue_mask = np.array([True] + [False] * (N - 1))
+        assert np.all(
+            issues["is_class_imbalance_issue"] == expected_issue_mask
+        ), "Issue mask should be correct"
+        expected_scores = np.array([0.01] + [1.0] * (N - 1))
+        np.testing.assert_allclose(
+            issues["class_imbalance_score"], expected_scores, err_msg="Scores should be correct"
+        )
+        assert summary["issue_type"][0] == "class_imbalance"
+        assert summary["score"][0] == pytest.approx(0.01, abs=0.01)
+
+    def test_report(self, create_issue_manager):
+        issue_manager = create_issue_manager()
+        issue_manager.find_issues()
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+        )
+        assert isinstance(report, str)
+        assert (
+            "------------------ class_imbalance issues ------------------\n\n"
+            "Number of examples with this issue:"
+        ) in report
+        assert ("Additional Information: \n" "Rarest Class:") in report
+
+    def test_collect_info(self, labels, create_issue_manager):
+        # With Imbalance
+        issue_manager = create_issue_manager(labels)
+        issue_manager.find_issues()
+        assert issue_manager.info["Rarest Class"] == 5
+
+        # Without Imbalance
+        labels[0] = 0
+        issue_manager = create_issue_manager(labels)
+        issue_manager.find_issues()
+        assert issue_manager.info["Rarest Class"] == "NA"

tests/datalab/issue_manager/test_knn_graph.py

@@ -0,0 +1,66 @@
+import math
+from hypothesis import HealthCheck, given, settings, strategies as st
+import pytest
+
+from .conftest import knn_graph_strategy
+
+
+class TestKNNGraph:
+    @staticmethod
+    def assert_distances_sorted(distances):
+        """Check distances are sorted in ascending order for each row."""
+        for row in distances:
+            assert all(row[i] <= row[i + 1] for i in range(len(row) - 1))
+
+    @staticmethod
+    def assert_indices_unique(indices):
+        """Check that neighbor indices are unique and don't have the row's index."""
+        for row_idx, row in enumerate(indices):
+            assert len(set(row)) == len(row)  # Check uniqueness
+            assert row_idx not in row  # Check that row's index is not in the row
+
+    @staticmethod
+    def assert_mutual_neighbors_have_same_distances(distances, indices):
+        """Verify that mutual neighbors have the same distances."""
+        for i in range(distances.shape[0]):
+            for j in indices[i]:
+                if i in indices[j]:
+                    d_ij = distances[i][list(indices[i]).index(j)]
+                    d_ji = distances[j][list(indices[j]).index(i)]
+                    assert math.isclose(
+                        d_ij, d_ji
+                    ), f"Distances between {i} and {j} do not match: {d_ij} vs {d_ji}"
+
+    @staticmethod
+    def assert_mutual_consistency_of_knn_distances(distances, indices):
+        """Verify the mutual consistency of k-NN distances:
+        For every point i and its neighbor j, ensure that the distance from i to j
+        cannot be smaller than the distance from any other neighbor k of j to j.
+        """
+        for i in range(distances.shape[0]):
+            for j in indices[i]:
+                d_ij = distances[i][list(indices[i]).index(j)]
+                j_neighbors_distances = distances[j]
+                if d_ij < max(j_neighbors_distances):
+                    assert (
+                        i in indices[j]
+                    ), f"Point {i} should be a neighbor of point {j}, it's closer than the farthest neighbor of {j}"
+
+    @pytest.mark.slow
+    @given(
+        knn_graph=knn_graph_strategy(
+            num_samples=st.integers(min_value=6, max_value=10),
+            k_neighbors=st.integers(min_value=2, max_value=5),
+        )
+    )
+    @settings(suppress_health_check=[HealthCheck.too_slow], max_examples=1000, deadline=None)
+    def test_knn_graph(self, knn_graph):
+        """Run through the property tests defined above."""
+        N = knn_graph.shape[0]
+        distances = knn_graph.data.reshape(N, -1)
+        indices = knn_graph.indices.reshape(N, -1)
+
+        self.assert_distances_sorted(distances)
+        self.assert_indices_unique(indices)
+        self.assert_mutual_neighbors_have_same_distances(distances, indices)
+        self.assert_mutual_consistency_of_knn_distances(distances, indices)

tests/datalab/issue_manager/test_knn_graph_helpers.py

@@ -0,0 +1,290 @@
+import numpy as np
+import pytest
+from scipy.sparse import csr_matrix
+from sklearn.neighbors import NearestNeighbors
+from cleanlab.datalab.internal.issue_manager.knn_graph_helpers import (
+    _process_knn_graph_from_inputs as _test_fn_1,  # Rename for testing purposes
+    num_neighbors_in_knn_graph as _get_num_neighbors,
+    set_knn_graph as _test_fn_2,  # Rename for testing purposes
+)
+from cleanlab.internal.neighbor.knn_graph import create_knn_graph_and_index as _make_knn
+
+
+class TestProcessKNNGraphFromInputs:
+
+    def test_knn_graph_provided_in_kwargs(self):
+        """
+        Test case 1: Verify that the function uses the knn_graph provided in kwargs.
+        """
+        kwargs = {"knn_graph": csr_matrix(np.random.random((10, 10)))}
+        statistics = {"weighted_knn_graph": None}
+        result = _test_fn_1(kwargs, statistics, k_for_recomputation=5)
+        assert isinstance(result, csr_matrix)
+        assert result.shape == (10, 10)
+
+    def test_knn_graph_stored_in_statistics(self):
+        """
+        Test case 2: Verify that the function uses the knn_graph stored in statistics
+        when not provided in kwargs.
+        """
+        kwargs = {"knn_graph": None}
+        statistics = {"weighted_knn_graph": csr_matrix(np.random.random((10, 10)))}
+        result = _test_fn_1(kwargs, statistics, k_for_recomputation=5)
+        assert isinstance(result, csr_matrix)
+        assert result.shape == (10, 10)
+
+    def test_knn_graph_precedence_of_kwargs_over_statistics(self):
+        """
+        Test case 3: Verify that the knn_graph provided in kwargs takes precedence
+        over the one stored in statistics. The knn_graph provided in kwargs is
+        ALWAYS used if available.
+        """
+        kwargs = {"knn_graph": csr_matrix(np.random.random((10, 10))), "k": 5}
+        statistics = {"weighted_knn_graph": csr_matrix(np.random.random((5, 5)))}
+        result = _test_fn_1(kwargs, statistics, k_for_recomputation=5)
+        assert isinstance(result, csr_matrix)
+        assert result.shape == (10, 10)
+
+        # Even if the statistics knn_graph is larger, the user-provided knn_graph is preferred
+        statistics = {"weighted_knn_graph": csr_matrix(np.random.random((15, 15)))}
+        result = _test_fn_1(kwargs, statistics, k_for_recomputation=5)
+        assert result.shape == (10, 10)
+
+        # Even if k is larger than the user-provided knn_graph, the user-provided knn_graph is preferred
+        kwargs = {"knn_graph": csr_matrix(np.random.random((10, 10)))}
+        statistics = {"weighted_knn_graph": csr_matrix(np.random.random((11, 11)))}
+        result = _test_fn_1(kwargs, statistics, k_for_recomputation=15)
+        assert result.shape == (10, 10)
+
+    def test_no_knn_graph_provided(self):
+        """
+        Test case 4: Verify that the function returns None when no knn_graph is provided
+        in either kwargs or statistics.
+        """
+        kwargs = {"knn_graph": None}
+        statistics = {"weighted_knn_graph": None}
+        result = _test_fn_1(kwargs, statistics, k_for_recomputation=0)
+        assert result is None
+
+    def test_insufficient_knn_graph(self):
+        """
+        Test case 5: Verify the behavior of the function when the knn_graph in
+        statistics is insufficient for the given k value.
+        """
+        k = 20
+        kwargs = {"knn_graph": csr_matrix(np.random.random((10, 10)))}
+        statistics = {"weighted_knn_graph": None}
+        result = _test_fn_1(kwargs, statistics, k)
+        assert result.shape == (10, 10)
+
+        kwargs = {"knn_graph": None}
+        statistics = {"weighted_knn_graph": csr_matrix(np.random.random((10, 10)))}
+        result = _test_fn_1(kwargs, statistics, k)
+        assert result is None
+
+        statistics = {"weighted_knn_graph": csr_matrix(np.random.random((21, 21)))}
+        result = _test_fn_1(kwargs, statistics, k)
+        assert result.shape == (21, 21)
+
+        # With sufficiently small k, the kwargs graph is preferred as it's explicitly provided by the user
+        k = 5
+        kwargs = {"knn_graph": csr_matrix(np.random.random((10, 10)))}
+        result = _test_fn_1(kwargs, statistics, k)
+        assert result.shape == (10, 10)
+
+        kwargs = {"knn_graph": None}
+        result = _test_fn_1(kwargs, statistics, k)
+        assert result.shape == (21, 21)
+
+
+class TestSetKNNGraph:
+
+    @pytest.fixture
+    def small_knn_graph(self):
+        knn_graph, _ = _make_knn(np.random.random((10, 5)), n_neighbors=5, metric="euclidean")
+        return knn_graph
+
+    @pytest.fixture
+    def mid_knn_graph(self):
+        knn_graph, _ = _make_knn(np.random.random((10, 5)), n_neighbors=7, metric="euclidean")
+        return knn_graph
+
+    @pytest.fixture
+    def large_knn_graph(self):
+        knn_graph, _ = _make_knn(np.random.random((10, 5)), n_neighbors=9, metric="euclidean")
+        return knn_graph
+
+    @pytest.fixture
+    def manhattan_knn_graph(self):
+        knn_graph, _ = _make_knn(np.random.random((10, 5)), n_neighbors=5, metric="manhattan")
+        return knn_graph
+
+    def test_knn_graph_provided_in_kwargs(self, small_knn_graph):
+        """
+        Test case 1: Verify that the function uses the knn_graph provided in kwargs.
+        """
+        features = np.random.random((10, 5))
+        find_issues_kwargs = {"knn_graph": small_knn_graph}
+        statistics = {"weighted_knn_graph": None}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=5, statistics=statistics
+        )
+        assert isinstance(result_graph, csr_matrix)
+        assert _get_num_neighbors(result_graph) == 5
+        assert result_metric == "euclidean"
+
+    def test_knn_graph_stored_in_statistics(self, small_knn_graph):
+        """
+        Test case 2: Verify that the function uses the knn_graph stored in statistics
+        when not provided in kwargs.
+        """
+        features = np.random.random((10, 5))
+        find_issues_kwargs = {"knn_graph": None}
+        statistics = {"weighted_knn_graph": small_knn_graph, "knn_metric": "euclidean"}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=5, statistics=statistics
+        )
+        assert isinstance(result_graph, csr_matrix)
+        assert _get_num_neighbors(result_graph) == 5
+        assert result_metric == "euclidean"
+
+        # Even if k is smaller than what is in statitics, the metric will cause a recompute
+        statistics = {"weighted_knn_graph": small_knn_graph, "knn_metric": "euclidean_outdated"}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=4, statistics=statistics
+        )
+        assert _get_num_neighbors(result_graph) == 4
+        assert result_metric == "euclidean"
+
+        # If the metric hasn't changed, but the value of k is larger than the stored knn_graph, the knn_graph is recomputed
+        statistics = {"weighted_knn_graph": small_knn_graph, "knn_metric": "euclidean"}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=6, statistics=statistics
+        )
+        assert _get_num_neighbors(result_graph) == 6
+        assert result_metric == "euclidean"
+
+    def test_knn_graph_precedence_of_kwargs_over_statistics(
+        self, small_knn_graph, mid_knn_graph, large_knn_graph
+    ):
+        """
+        Test case 3: Verify that the knn_graph provided in kwargs takes precedence
+        over the one stored in statistics. The knn_graph provided in kwargs is
+        ALWAYS used if available.
+        """
+        features = np.random.random((10, 5))
+        find_issues_kwargs = {"knn_graph": mid_knn_graph}
+        statistics = {"weighted_knn_graph": small_knn_graph, "knn_metric": "euclidean"}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=5, statistics=statistics
+        )
+        assert isinstance(result_graph, csr_matrix)
+        assert _get_num_neighbors(result_graph) == 7
+        assert result_metric == "euclidean"
+
+        # Even if the statistics knn_graph is larger, the user-provided knn_graph is preferred
+        statistics = {"weighted_knn_graph": large_knn_graph, "knn_metric": "euclidean"}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=5, statistics=statistics
+        )
+        assert _get_num_neighbors(result_graph) == 7
+        assert result_metric == "euclidean"
+
+        # Even if k is larger than the user-provided knn_graph, the user-provided knn_graph is preferred
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=8, statistics=statistics
+        )
+        assert _get_num_neighbors(result_graph) == 7
+        assert result_metric == "euclidean"
+
+    def test_no_knn_graph_provided(self):
+        """
+        Test case 4: Verify that the function creates a new knn_graph when no knn_graph
+        is provided in either kwargs or statistics. Features are required.
+        """
+        features = np.random.random((10, 5))
+        find_issues_kwargs = {"knn_graph": None}
+        statistics = {"weighted_knn_graph": None}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="cosine", k=3, statistics=statistics
+        )
+        assert _get_num_neighbors(result_graph) == 3
+        assert result_metric == "cosine"
+
+        with pytest.raises(
+            AssertionError, match="Features must be provided to compute the knn graph."
+        ):
+            _test_fn_2(None, find_issues_kwargs, metric="cosine", k=0, statistics=statistics)
+
+    def test_metric_change_requires_new_knn_graph(self, manhattan_knn_graph):
+        """
+        Test case 5: Verify that the function creates a new knn_graph if the metric has changed.
+        """
+        features = np.random.random((10, 5))
+        find_issues_kwargs = {"knn_graph": None}
+        statistics = {"weighted_knn_graph": manhattan_knn_graph, "knn_metric": "manhattan"}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=2, statistics=statistics
+        )
+        assert _get_num_neighbors(result_graph) == 2
+        assert result_metric == "euclidean"
+
+    def test_knn_graph_with_insufficient_graph(self, small_knn_graph, large_knn_graph):
+        """
+        Test case 6: Verify the behavior of the function when the knn_graph in
+        statistics is insufficient for the given k value.
+        """
+        features = np.random.random((10, 5))
+        k = 8
+        find_issues_kwargs = {"knn_graph": small_knn_graph}
+        statistics = {"weighted_knn_graph": None}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=k, statistics=statistics
+        )
+        assert _get_num_neighbors(result_graph) == 5
+        assert result_metric == "euclidean"
+
+        # The small graph doesn't have enough neighbors, so it should be recomputed
+        find_issues_kwargs = {"knn_graph": None}
+        statistics = {"weighted_knn_graph": small_knn_graph}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=k, statistics=statistics
+        )
+        assert _get_num_neighbors(result_graph) == 8
+        assert result_metric is "euclidean"
+
+        # The large graph has more than enough neighbors, so it should be used
+        statistics = {"weighted_knn_graph": large_knn_graph}
+        result_graph, result_metric, _ = _test_fn_2(
+            features, find_issues_kwargs, metric="euclidean", k=k, statistics=statistics
+        )
+        assert _get_num_neighbors(result_graph) == 9
+        assert result_metric is "euclidean"
+
+    def test_knn_returned(self, small_knn_graph):
+        features = np.random.random((10, 5))
+        k = 3
+        result_graph, result_metric, result_knn = _test_fn_2(
+            features, {"knn_graph": None}, metric="cosine", k=k, statistics={}
+        )
+        assert isinstance(result_knn, NearestNeighbors)
+        assert result_knn.n_neighbors == k
+        assert result_knn.metric == "cosine"
+
+        result_graph, result_metric, result_knn = _test_fn_2(
+            features, {"knn_graph": small_knn_graph}, metric="euclidean", k=k, statistics={}
+        )
+        assert result_knn == None
+        assert result_metric == "euclidean"
+        np.testing.assert_array_equal(result_graph.toarray(), small_knn_graph.toarray())
+
+        result_graph, result_metric, result_knn = _test_fn_2(
+            features,
+            {"knn_graph": None},
+            metric="euclidean",
+            k=k,
+            statistics={"weighted_knn_graph": small_knn_graph, "knn_metric": "euclidean"},
+        )
+        assert result_knn == None
+        assert result_metric == "euclidean"
+        np.testing.assert_array_equal(result_graph.toarray(), small_knn_graph.toarray())

tests/datalab/issue_manager/test_label.py

@@ -0,0 +1,116 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from cleanlab.datalab.internal.issue_manager.label import LabelIssueManager
+
+
+class TestLabelIssueManager:
+    @pytest.fixture
+    def issue_manager(self, lab):
+        return LabelIssueManager(datalab=lab)
+
+    def test_find_issues(self, lab, pred_probs, issue_manager):
+        """Test that the find_issues method works."""
+        issue_manager.find_issues(pred_probs=pred_probs)
+        issues, summary, info = issue_manager.issues, issue_manager.summary, issue_manager.info
+        assert isinstance(issues, pd.DataFrame), "Issues should be a dataframe"
+
+        assert isinstance(summary, pd.DataFrame), "Summary should be a dataframe"
+        assert summary["issue_type"].values[0] == "label"
+        assert pytest.approx(summary["score"].values[0]) == 0.4
+
+        assert isinstance(info, dict), "Info should be a dict"
+        info_keys = info.keys()
+        expected_keys = [
+            "num_label_issues",
+            "average_label_quality",
+            "confident_joint",
+            "classes_by_label_quality",
+            "overlapping_classes",
+            "py",
+            "noise_matrix",
+            "inverse_noise_matrix",
+        ]
+        assert all(
+            [key in info_keys for key in expected_keys]
+        ), f"Info should have the right keys, but is missing {set(expected_keys) - set(info_keys)}"
+        # Compare results with low_memory=True
+        clean_learning_kwargs = {"low_memory": True}
+        issue_manager_lm = LabelIssueManager(
+            datalab=lab, clean_learning_kwargs=clean_learning_kwargs
+        )
+        issue_manager_lm.find_issues(pred_probs=pred_probs)
+        issues_lm = issue_manager_lm.issues
+        # jaccard similarity
+        intersection = len(list(set(issues).intersection(set(issues_lm))))
+        union = len(set(issues)) + len(set(issues_lm)) - intersection
+        assert float(intersection) / union > 0.95
+
+    def test_find_issues_with_kwargs(self, pred_probs, issue_manager):
+        issue_manager.find_issues(pred_probs=pred_probs, thresholds=[0.2, 0.3, 0.1])
+
+    def test_init_with_clean_learning_kwargs(self, lab, issue_manager):
+        """Test that the init method can provide kwargs to the CleanLearning constructor."""
+        new_issue_manager = LabelIssueManager(
+            datalab=lab,
+            clean_learning_kwargs={"cv_n_folds": 10},
+        )
+        cv_n_folds = [im.cl.cv_n_folds for im in [issue_manager, new_issue_manager]]
+        assert cv_n_folds == [5, 10], "Issue manager should have the right attributes"
+
+    def test_get_summary_parameters(self, issue_manager, monkeypatch):
+        mock_health_summary_parameters = {
+            "labels": [1, 0, 2],
+            "asymmetric": False,
+            "class_names": ["a", "b", "c"],
+            "num_examples": 3,
+            "joint": [1 / 3, 1 / 3, 1 / 3],
+            "confident_joint": [1 / 3, 1 / 3, 1 / 3],
+            "multi_label": False,
+        }
+        pred_probs = np.random.rand(3, 3)
+        monkeypatch.setattr(
+            issue_manager, "health_summary_parameters", mock_health_summary_parameters
+        )
+        summary_parameters = issue_manager._get_summary_parameters(pred_probs=pred_probs)
+        expected_parameters = {
+            "confident_joint": [1 / 3, 1 / 3, 1 / 3],
+            "asymmetric": False,
+            "class_names": ["a", "b", "c"],
+        }
+        assert summary_parameters == expected_parameters
+
+        # Test missing "confident_joint" key
+        mock_health_summary_parameters.pop("confident_joint")
+        monkeypatch.setattr(
+            issue_manager, "health_summary_parameters", mock_health_summary_parameters
+        )
+        summary_parameters = issue_manager._get_summary_parameters(pred_probs=pred_probs)
+        expected_parameters = {
+            "joint": [1 / 3, 1 / 3, 1 / 3],
+            "num_examples": 3,
+            "asymmetric": False,
+            "class_names": ["a", "b", "c"],
+        }
+        assert summary_parameters == expected_parameters
+
+        # Test missing "joint" key
+        mock_health_summary_parameters.pop("joint")
+        monkeypatch.setattr(
+            issue_manager.datalab._labels, "labels", mock_health_summary_parameters["labels"]
+        )
+        monkeypatch.setattr(
+            issue_manager, "health_summary_parameters", mock_health_summary_parameters
+        )
+        summary_parameters = issue_manager._get_summary_parameters(pred_probs=pred_probs)
+        expected_parameters = {
+            "pred_probs": pred_probs,
+            "labels": [1, 0, 2],
+            "asymmetric": False,
+            "class_names": ["a", "b", "c"],
+        }
+        assert np.all(summary_parameters["pred_probs"] == expected_parameters["pred_probs"])
+        summary_parameters.pop("pred_probs")
+        expected_parameters.pop("pred_probs")
+        assert summary_parameters == expected_parameters

tests/datalab/issue_manager/test_noniid.py

@@ -0,0 +1,339 @@
+import numpy as np
+import pytest
+
+from cleanlab.datalab.internal.issue_manager.noniid import (
+    NonIIDIssueManager,
+    simplified_kolmogorov_smirnov_test,
+)
+
+SEED = 42
+
+
+@pytest.mark.parametrize(
+    "neighbor_histogram, non_neighbor_histogram, expected_statistic",
+    [
+        # Test with equal histograms
+        (
+            [0.25, 0.25, 0.25, 0.25],
+            [0.25, 0.25, 0.25, 0.25],
+            0.0,
+        ),
+        # Test with maximum difference in the first bin
+        (
+            [1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.25, 0.25, 0.5],
+            1.0,
+        ),
+        # Test with maximum difference in the last bin
+        (
+            [0.25, 0.25, 0.25, 0.25],
+            [0.5, 0.25, 0.25, 0.0],
+            0.25,
+        ),
+        # Test with arbitrary histograms
+        (
+            [0.2, 0.3, 0.4, 0.1],
+            [0.1, 0.4, 0.25, 0.3],
+            0.15,  # (0.2 -> 0.5 -> *0.9* -> 1.0) vs (0.1 -> 0.5 -> *0.75* -> 1.05
+        ),
+    ],
+    ids=[
+        "equal_histograms",
+        "maximum_difference_in_first_bin",
+        "maximum_difference_in_last_bin",
+        "arbitrary_histograms",
+    ],
+)
+def test_simplified_kolmogorov_smirnov_test(
+    neighbor_histogram, non_neighbor_histogram, expected_statistic
+):
+    nh = np.array(neighbor_histogram)
+    nnh = np.array(non_neighbor_histogram)
+    statistic = simplified_kolmogorov_smirnov_test(nh, nnh)
+    np.testing.assert_almost_equal(statistic, expected_statistic)
+
+
+class TestNonIIDIssueManager:
+    @pytest.fixture
+    def embeddings(self, lab):
+        np.random.seed(SEED)
+        embeddings_array = np.arange(lab.get_info("statistics")["num_examples"] * 10).reshape(-1, 1)
+        return embeddings_array
+
+    @pytest.fixture
+    def pred_probs(self, lab):
+        pred_probs_array = (
+            np.arange(lab.get_info("statistics")["num_examples"] * 10).reshape(-1, 1)
+        ) / len(np.arange(lab.get_info("statistics")["num_examples"] * 10).reshape(-1, 1))
+        return pred_probs_array
+
+    @pytest.fixture
+    def issue_manager(self, lab):
+        return NonIIDIssueManager(
+            datalab=lab,
+            metric="euclidean",
+            k=10,
+        )
+
+    def test_init(self, lab, issue_manager):
+        assert issue_manager.datalab == lab
+        assert issue_manager.metric == "euclidean"
+        assert issue_manager.k == 10
+        assert issue_manager.num_permutations == 25
+        assert issue_manager.significance_threshold == 0.05
+
+        issue_manager = NonIIDIssueManager(
+            datalab=lab,
+            num_permutations=15,
+        )
+
+        assert issue_manager.num_permutations == 15
+
+    def test_find_issues(self, issue_manager, embeddings):
+        np.random.seed(SEED)
+        issue_manager.find_issues(features=embeddings)
+        issues_sort, summary_sort, info_sort = (
+            issue_manager.issues,
+            issue_manager.summary,
+            issue_manager.info,
+        )
+        expected_sorted_issue_mask = np.array([False] * 46 + [True] + [False] * 3)
+        assert np.all(
+            issues_sort["is_non_iid_issue"] == expected_sorted_issue_mask
+        ), "Issue mask should be correct"
+        assert summary_sort["issue_type"][0] == "non_iid"
+        assert summary_sort["score"][0] == pytest.approx(expected=0.0, abs=1e-7)
+        assert info_sort.get("p-value", None) is not None, "Should have p-value"
+        assert summary_sort["score"][0] == pytest.approx(expected=info_sort["p-value"], abs=1e-7)
+
+        permutation = np.random.permutation(len(embeddings))
+        new_issue_manager = NonIIDIssueManager(
+            datalab=issue_manager.datalab,
+            metric="euclidean",
+            k=10,
+        )
+        new_issue_manager.find_issues(features=embeddings[permutation])
+        issues_perm, summary_perm, info_perm = (
+            new_issue_manager.issues,
+            new_issue_manager.summary,
+            new_issue_manager.info,
+        )
+        expected_permuted_issue_mask = np.array([False] * len(embeddings))
+        assert np.all(
+            issues_perm["is_non_iid_issue"] == expected_permuted_issue_mask
+        ), "Issue mask should be correct"
+        assert summary_perm["issue_type"][0] == "non_iid"
+        # ensure score is large, cannot easily ensure precise value because random seed has different effects on
+        # different OS:
+        assert summary_perm["score"][0] > 0.05
+        assert info_perm.get("p-value", None) is not None, "Should have p-value"
+        assert summary_perm["score"][0] == pytest.approx(expected=info_perm["p-value"], abs=1e-7)
+
+    def test_find_issues_using_pred_probs(self, issue_manager, pred_probs):
+        np.random.seed(SEED)
+        issue_manager.find_issues(pred_probs=pred_probs)
+        issues_sort, summary_sort, info_sort = (
+            issue_manager.issues,
+            issue_manager.summary,
+            issue_manager.info,
+        )
+        expected_sorted_issue_mask = np.array([False] * 46 + [True] + [False] * 3)
+        assert np.all(
+            issues_sort["is_non_iid_issue"] == expected_sorted_issue_mask
+        ), "Issue mask should be correct"
+        assert summary_sort["issue_type"][0] == "non_iid"
+        assert summary_sort["score"][0] == pytest.approx(expected=0.0, abs=1e-7)
+        assert info_sort.get("p-value", None) is not None, "Should have p-value"
+        assert summary_sort["score"][0] == pytest.approx(expected=info_sort["p-value"], abs=1e-7)
+
+        permutation = np.random.permutation(len(pred_probs))
+        new_issue_manager = NonIIDIssueManager(
+            datalab=issue_manager.datalab,
+            metric="euclidean",
+            k=10,
+        )
+
+        new_issue_manager.find_issues(pred_probs=pred_probs[permutation])
+        issues_perm, summary_perm, info_perm = (
+            new_issue_manager.issues,
+            new_issue_manager.summary,
+            new_issue_manager.info,
+        )
+        expected_permuted_issue_mask = np.array([False] * len(pred_probs))
+        assert np.all(
+            issues_perm["is_non_iid_issue"] == expected_permuted_issue_mask
+        ), "Issue mask should be correct"
+        assert summary_perm["issue_type"][0] == "non_iid"
+        # ensure score is large, cannot easily ensure precise value because random seed has different effects on
+        # different OS:
+        assert summary_perm["score"][0] > 0.05
+        assert info_perm.get("p-value", None) is not None, "Should have p-value"
+        assert summary_perm["score"][0] == pytest.approx(expected=info_perm["p-value"], abs=1e-7)
+
+    def test_report(self, issue_manager, embeddings):
+        np.random.seed(SEED)
+        issue_manager.find_issues(features=embeddings)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+        )
+
+        assert isinstance(report, str)
+        assert (
+            "---------------------- non_iid issues ----------------------\n\n"
+            "Number of examples with this issue:"
+        ) in report
+
+        issue_manager.find_issues(features=embeddings)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+            verbosity=3,
+        )
+
+        assert "Additional Information: " in report
+
+    def test_report_using_pred_probs(self, issue_manager, pred_probs):
+        np.random.seed(SEED)
+        issue_manager.find_issues(pred_probs=pred_probs)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+        )
+
+        assert (
+            "---------------------- non_iid issues ----------------------\n\n"
+            "Number of examples with this issue:"
+        ) in report
+
+        issue_manager.find_issues(pred_probs=pred_probs)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+            verbosity=3,
+        )
+
+        assert "Additional Information: " in report
+
+    def test_collect_info(self, issue_manager, embeddings):
+        """Test some values in the info dict.
+
+        Mainly focused on the nearest neighbor info.
+        """
+
+        issue_manager.find_issues(features=embeddings)
+        info = issue_manager.info
+
+        assert info["p-value"] == 0
+        assert info["metric"] == "euclidean"
+        assert info["k"] == 10
+
+    def test_collect_info_using_pred_probs(self, issue_manager, pred_probs):
+        """Test some values in the info dict.
+
+        Mainly focused on the nearest neighbor info.
+        """
+        issue_manager.find_issues(pred_probs=pred_probs)
+        info = issue_manager.info
+
+        assert info["p-value"] == 0
+        assert info["metric"] == "euclidean"
+        assert info["k"] == 10
+
+    @pytest.mark.parametrize(
+        "seed",
+        [
+            "default",
+            SEED,
+            None,
+        ],
+        ids=["default", "seed", "no_seed"],
+    )
+    def test_seed(self, lab, seed):
+        num_classes = 10
+        means = [
+            np.array([np.random.uniform(high=10), np.random.uniform(high=10)])
+            for _ in range(num_classes)
+        ]
+        sigmas = [np.random.uniform(high=1) for _ in range(num_classes)]
+        class_stats = list(zip(means, sigmas))
+        num_samples = 2000
+
+        def generate_data_iid():
+            # This should be IID, resulting in a larger p-value
+            samples = []
+            labels = []
+
+            for _ in range(num_samples):
+                label = np.random.choice(num_classes)
+                mean, sigma = class_stats[label]
+                sample = np.random.normal(mean, sigma)
+                samples.append(sample)
+                labels.append(label)
+            samples = np.array(samples)
+            labels = np.array(labels)
+            dataset = {"features": samples, "labels": labels}
+            return dataset
+
+        dataset = generate_data_iid()
+        embeddings = dataset["features"]
+
+        # Create new issue manager, ignore the lab assigned for this test
+        if seed == "default":
+            issue_manager = NonIIDIssueManager(
+                datalab=lab,
+                metric="euclidean",
+                k=10,
+            )
+        else:
+            issue_manager = NonIIDIssueManager(
+                datalab=lab,
+                metric="euclidean",
+                k=10,
+                seed=seed,
+            )
+        issue_manager.find_issues(features=embeddings)
+        p_value = issue_manager.info["p-value"]
+
+        # Run again with the same seed
+        issue_manager.find_issues(features=embeddings)
+        p_value2 = issue_manager.info["p-value"]
+
+        assert p_value > 0.0
+        if seed is not None or seed == "default":
+            assert p_value == p_value2
+        else:
+            assert p_value != p_value2
+
+        # using pred_probs
+        # normalizing pred_probs (0 to 1)
+        pred_probs = embeddings / (np.max(embeddings) - np.min(embeddings))
+        if seed == "default":
+            issue_manager = NonIIDIssueManager(
+                datalab=lab,
+                metric="euclidean",
+                k=10,
+            )
+        else:
+            issue_manager = NonIIDIssueManager(
+                datalab=lab,
+                metric="euclidean",
+                k=10,
+                seed=seed,
+            )
+        issue_manager.find_issues(pred_probs=pred_probs)
+        p_value = issue_manager.info["p-value"]
+
+        # Run again with the same seed
+        issue_manager.find_issues(pred_probs=pred_probs)
+        p_value2 = issue_manager.info["p-value"]
+
+        assert p_value > 0.0
+        if seed is not None or seed == "default":
+            assert p_value == p_value2
+        else:
+            assert p_value != p_value2

tests/datalab/issue_manager/test_null.py

@@ -0,0 +1,196 @@
+import numpy as np
+import pandas as pd
+import pytest
+from hypothesis import HealthCheck, given, settings
+from hypothesis.extra.numpy import array_shapes, arrays
+from hypothesis.strategies import floats, just
+
+from cleanlab.datalab.internal.issue_manager.null import NullIssueManager
+
+SEED = 42
+
+
+class TestNullIssueManager:
+    @pytest.fixture
+    def embeddings(self):
+        np.random.seed(SEED)
+        embeddings_array = np.random.random((4, 3))
+        return embeddings_array
+
+    @pytest.fixture
+    def embeddings_with_null(self):
+        np.random.seed(SEED)
+        embeddings_array = np.random.random((4, 3))
+        embeddings_array[[0, 3], 0] = np.nan
+        embeddings_array[1] = np.nan
+        return embeddings_array
+
+    @pytest.fixture
+    def issue_manager(self, lab):
+        return NullIssueManager(datalab=lab)
+
+    def test_init(self, lab, issue_manager):
+        assert issue_manager.datalab == lab
+
+    def test_find_issues(self, issue_manager, embeddings):
+        np.random.seed(SEED)
+        issue_manager.find_issues(features=embeddings)
+        issues_sort, summary_sort, info_sort = (
+            issue_manager.issues,
+            issue_manager.summary,
+            issue_manager.info,
+        )
+        expected_sorted_issue_mask = np.array([False, False, False, False])
+        assert np.all(
+            issues_sort["is_null_issue"] == expected_sorted_issue_mask
+        ), "Issue mask should be correct"
+        assert summary_sort["issue_type"][0] == "null"
+        assert summary_sort["score"][0] == pytest.approx(expected=1.0, abs=1e-7)
+        assert (
+            info_sort.get("average_null_score", None) is not None
+        ), "Should have average null score"
+        assert summary_sort["score"][0] == pytest.approx(
+            expected=info_sort["average_null_score"], abs=1e-7
+        )
+
+    def test_find_issues_with_null(self, issue_manager, embeddings_with_null):
+        np.random.seed(SEED)
+        issue_manager.find_issues(features=embeddings_with_null)
+        issues_sort, summary_sort, info_sort = (
+            issue_manager.issues,
+            issue_manager.summary,
+            issue_manager.info,
+        )
+        expected_sorted_issue_mask = np.array([False, True, False, False])
+        assert np.all(
+            issues_sort["is_null_issue"] == expected_sorted_issue_mask
+        ), "Issue mask should be correct"
+        assert summary_sort["issue_type"][0] == "null"
+        assert summary_sort["score"][0] == pytest.approx(expected=7 / 12, abs=1e-7)
+        assert (
+            info_sort.get("average_null_score", None) is not None
+        ), "Should have average null score"
+        assert summary_sort["score"][0] == pytest.approx(
+            expected=info_sort["average_null_score"], abs=1e-7
+        )
+
+    def test_report(self, issue_manager, embeddings):
+        np.random.seed(SEED)
+        issue_manager.find_issues(features=embeddings)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+        )
+
+        assert isinstance(report, str)
+        assert (
+            "----------------------- null issues ------------------------\n\n"
+            "Number of examples with this issue:"
+        ) in report
+
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+            verbosity=3,
+        )
+        assert "Additional Information: " in report
+
+    def test_report_with_null(self, issue_manager, embeddings_with_null):
+        np.random.seed(SEED)
+        issue_manager.find_issues(features=embeddings_with_null)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+        )
+
+        assert isinstance(report, str)
+        assert (
+            "----------------------- null issues ------------------------\n\n"
+            "Number of examples with this issue:"
+        ) in report
+
+        assert "Additional Information: " not in report
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+            verbosity=3,
+        )
+        assert "Additional Information: " in report
+
+    def test_collect_info(self, issue_manager, embeddings):
+        """Test some values in the info dict."""
+        issue_manager.find_issues(features=embeddings)
+        info = issue_manager.info
+        assert info["average_null_score"] == pytest.approx(1.0, abs=0.01)
+        assert info["most_common_issue"]["pattern"] == "no_null"
+        assert info["most_common_issue"]["count"] == 0
+        assert info["most_common_issue"]["rows_affected"] == []
+        assert info["column_impact"] == [0, 0, 0]
+
+    def test_collect_info_with_nulls(self, issue_manager, embeddings_with_null):
+        """Test some values in the info dict."""
+        issue_manager.find_issues(features=embeddings_with_null)
+        info = issue_manager.info
+        assert info["average_null_score"] == pytest.approx(expected=7 / 12, abs=1e-7)
+        assert info["most_common_issue"]["pattern"] == "100"
+        assert info["most_common_issue"]["count"] == 2
+        assert info["most_common_issue"]["rows_affected"] == [0, 3]
+        assert info["column_impact"] == [0.75, 0.25, 0.25]
+
+    def test_can_work_with_different_dtypes(self, issue_manager):
+        features = pd.DataFrame(
+            {
+                "bool": [True, False, True, False],
+                "object": [True, False, True, np.nan],
+                "uint8": np.array([0, 1, 3, 4], dtype=np.uint8),
+                "int8": np.array([0, -1, 3, -4], dtype=np.int8),
+                "float": [0.1, np.nan, 0.3, 0.4],
+            }
+        )
+        issue_manager.find_issues(features=features)
+        info = issue_manager.info
+        assert info["average_null_score"] == pytest.approx(expected=18 / 20, abs=1e-7)
+        assert info["column_impact"] == [0, 0.25, 0, 0, 0.25]
+
+    # Strategy for generating NaN values
+    nan_strategy = just(np.nan)
+
+    # Strategy for generating regular float values, including NaNs
+    float_with_nan = floats(allow_nan=True)
+
+    # Strategy for generating NumPy arrays with some NaN values
+    features_with_nan_strategy = arrays(
+        dtype=np.float64,
+        shape=array_shapes(min_dims=2, max_dims=2, min_side=1, max_side=5),
+        elements=float_with_nan,
+        fill=nan_strategy,
+    )
+
+    @settings(
+        suppress_health_check=[HealthCheck.function_scoped_fixture],
+        deadline=None,
+    )  # No need to reset state of issue_manager fixture
+    @given(embeddings=features_with_nan_strategy)
+    def test_quality_scores_and_full_null_row_identification(self, issue_manager, embeddings):
+        # Run the find_issues method
+        issue_manager.find_issues(features=embeddings)
+        issues_sort, _, _ = (
+            issue_manager.issues,
+            issue_manager.summary,
+            issue_manager.info,
+        )
+
+        # Check for the two main properties:
+
+        # 1. The quality score for each row should be the fraction of features which are not null in that row.
+        non_null_fractions = [np.count_nonzero(~np.isnan(row)) / len(row) for row in embeddings]
+        scores = issues_sort[issue_manager.issue_score_key]
+        assert np.allclose(scores, non_null_fractions, atol=1e-7)
+
+        # 2. The rows that are marked as is_null_issue should ONLY be those rows which are 100% null values.
+        all_rows_are_null = np.all(np.isnan(embeddings), axis=1)
+        assert np.all(issues_sort["is_null_issue"] == all_rows_are_null)

tests/datalab/issue_manager/test_outlier.py

@@ -0,0 +1,190 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from cleanlab.datalab.internal.issue_manager.outlier import OutlierIssueManager
+from cleanlab.outlier import OutOfDistribution
+
+SEED = 42
+
+
+class TestOutlierIssueManager:
+    @pytest.fixture
+    def embeddings(self, lab):
+        np.random.seed(SEED)
+        embeddings_array = 0.5 + 0.1 * np.random.rand(lab.get_info("statistics")["num_examples"], 2)
+        embeddings_array[4, :] = -1
+        return {"embedding": embeddings_array}
+
+    @pytest.fixture
+    def issue_manager(self, lab):
+        return OutlierIssueManager(datalab=lab, k=3)
+
+    @pytest.fixture
+    def issue_manager_with_threshold(self, lab):
+        return OutlierIssueManager(datalab=lab, k=2, threshold=0.5)
+
+    def test_init(self, issue_manager, issue_manager_with_threshold):
+        assert isinstance(issue_manager.ood, OutOfDistribution)
+        assert issue_manager.ood.params["k"] == 3
+        assert issue_manager.threshold == None
+
+        assert issue_manager_with_threshold.ood.params["k"] == 2
+        assert issue_manager_with_threshold.threshold == 0.5
+
+    def test_find_issues(self, issue_manager, issue_manager_with_threshold, embeddings):
+        issue_manager.find_issues(features=embeddings["embedding"])
+        issues, summary, info = issue_manager.issues, issue_manager.summary, issue_manager.info
+        expected_issue_mask = np.array([False] * 4 + [True])
+        assert np.all(
+            issues["is_outlier_issue"] == expected_issue_mask
+        ), "Issue mask should be correct"
+        # Assert that the argsort is correct
+        assert np.all(
+            issues["outlier_score"].argsort() == np.array([4, 2, 1, 3, 0])
+        ), "Outlier scores should be correct"
+
+        assert summary["issue_type"][0] == "outlier"
+        assert summary["score"][0] == pytest.approx(expected=0.3028243, abs=1e-7)
+
+        # New test data points are considered outliers if their average knn distance is greater than this issue threshold.
+        assert info.get("issue_threshold", None) is not None, "Should have issue_threshold info"
+        assert info.get("ood", None) is not None, "Should have the OutOfDistribution object in info"
+        assert issue_manager.threshold == pytest.approx(expected=0.37037, abs=1e-5)
+
+        issue_manager_with_threshold.find_issues(features=embeddings["embedding"])
+
+    def test_find_issues_with_pred_probs(self, lab):
+        issue_manager = OutlierIssueManager(datalab=lab, threshold=0.3)
+        pred_probs = np.array(
+            [
+                [0.25, 0.725, 0.025],
+                [0.37, 0.42, 0.21],
+                [0.05, 0.05, 0.9],
+                [0.1, 0.05, 0.85],
+                [0.1125, 0.65, 0.2375],
+            ]
+        )
+        issue_manager.find_issues(pred_probs=pred_probs)
+        issues, summary, info = issue_manager.issues, issue_manager.summary, issue_manager.info
+        expected_issue_mask = np.array([False] * 4 + [True])
+        assert np.all(
+            issues["is_outlier_issue"] == expected_issue_mask
+        ), "Issue mask should be correct"
+        assert summary["issue_type"][0] == "outlier"
+        assert summary["score"][0] == pytest.approx(expected=0.210, abs=1e-3)
+
+        assert issue_manager.threshold == 0.3
+
+        assert np.all(
+            info.get("confident_thresholds", None) == [0.1, 0.5725, 0.56875]
+        ), "Should have confident_joint info"
+
+    def test_find_issues_with_different_thresholds(self, lab, embeddings):
+        issue_manager = OutlierIssueManager(datalab=lab, k=3, threshold=0.66666)
+        issue_manager.find_issues(features=embeddings["embedding"])
+        issues, summary, info = issue_manager.issues, issue_manager.summary, issue_manager.info
+        expected_issue_mask = np.array([False] * 4 + [True])
+        assert np.all(
+            issues["is_outlier_issue"] == expected_issue_mask
+        ), "Issue mask should be correct"
+        # Assert that the argsort is correct
+        assert np.all(
+            issues["outlier_score"].argsort() == np.array([4, 2, 1, 3, 0])
+        ), "Outlier scores should be correct"
+
+        assert summary["issue_type"][0] == "outlier"
+        assert summary["score"][0] == pytest.approx(expected=0.3028243, abs=1e-7)
+
+        assert issue_manager.threshold == pytest.approx(0.66666, abs=0.01)
+
+    def test_report(self, issue_manager):
+        pred_probs = np.array(
+            [
+                [0.1, 0.85, 0.05],
+                [0.15, 0.8, 0.05],
+                [0.05, 0.05, 0.9],
+                [0.1, 0.05, 0.85],
+                [0.1, 0.65, 0.25],
+            ]
+        )
+        issue_manager.find_issues(pred_probs=pred_probs)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+        )
+        assert isinstance(report, str)
+        assert (
+            "---------------------- outlier issues ----------------------\n\n"
+            "Number of examples with this issue:"
+        ) in report
+
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+            verbosity=3,
+        )
+        assert "Additional Information: " in report
+
+        # Mock some vector and matrix values in the info dict
+        mock_info = issue_manager.info
+        vector = np.array([1, 2, 3, 4, 5, 6])
+        matrix = np.array([[i for i in range(20)] for _ in range(10)])
+        df = pd.DataFrame(matrix)
+        mock_list = [9, 8, 7, 6, 5, 4, 3, 2, 1]
+        mock_dict = {"a": 1, "b": 2, "c": 3}
+        mock_info["vector"] = vector
+        mock_info["matrix"] = matrix
+        mock_info["list"] = mock_list
+        mock_info["dict"] = mock_dict
+        mock_info["df"] = df
+
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info={**issue_manager.info, **mock_info},
+            verbosity=4,
+        )
+        assert "Additional Information: " in report
+        assert "vector: [1, 2, 3, 4, '...']" in report
+        assert f"matrix: array of shape {matrix.shape}\n[[ 0 " in report
+        assert "list: [9, 8, 7, 6, '...']" in report
+        assert 'dict:\n{\n    "a": 1,\n    "b": 2,\n    "c": 3\n}' in report
+        assert "df:" in report
+
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info={**issue_manager.info, **mock_info},
+            verbosity=2,
+        )
+        assert "Additional Information: " in report
+        assert "vector: [1, 2, 3, 4, '...']" not in report
+        assert f"matrix: array of shape {matrix.shape}\n[[ 0 " not in report
+        assert "list: [9, 8, 7, 6, '...']" not in report
+        assert 'dict:\n{\n    "a": 1,\n    "b": 2,\n    "c": 3\n}' not in report
+        assert "df:" not in report
+
+    def test_collect_info(self, issue_manager, embeddings):
+        """Test some values in the info dict.
+
+        Mainly focused on the nearest neighbor info.
+        """
+
+        issue_manager.find_issues(features=embeddings["embedding"])
+        info = issue_manager.info
+
+        nearest_neighbors = info["nearest_neighbor"]
+        distances_to_nearest_neighbor = info["distance_to_nearest_neighbor"]
+
+        assert nearest_neighbors == [3, 0, 3, 0, 2], "Nearest neighbors should be correct"
+
+        assert pytest.approx(distances_to_nearest_neighbor, abs=1e-3) == [
+            0.033,
+            0.05,
+            0.072,
+            0.033,
+            2.143,
+        ], "Distances to nearest neighbor should be correct"

tests/datalab/issue_manager/test_underperforming_group.py

@@ -0,0 +1,240 @@
+import numpy as np
+import pytest
+import pandas as pd
+import scipy.sparse as sp
+
+from cleanlab.datalab.internal.issue_manager.underperforming_group import (
+    UnderperformingGroupIssueManager,
+)
+from sklearn.datasets import make_blobs, load_iris
+
+SEED = 42
+
+
+class TestUnderperformingGroupIssueManager:
+    def generate_pred_probs(self, N, K, labels, noisy=False):
+        pred_probs = np.full((N, K), 0.1)
+        pred_probs[np.arange(N), labels] = 0.9
+        pred_probs = pred_probs / np.sum(pred_probs, axis=-1, keepdims=True)
+        if noisy:  # Swap columns of a class to generate incorrect predictions
+            pred_probs_slice = pred_probs[labels == 0]
+            pred_probs_slice[:, [0, 1]] = pred_probs_slice[:, [1, 0]]
+            pred_probs[labels == 0] = pred_probs_slice
+        return pred_probs
+
+    @pytest.fixture
+    def make_data(self, noisy=False):
+        def data(noisy=noisy):
+            N = 400
+            K = 4
+            features, labels = make_blobs(n_samples=N, centers=K, n_features=2, random_state=SEED)
+            pred_probs = self.generate_pred_probs(N, K, labels, noisy)
+            data = {"features": features, "pred_probs": pred_probs, "labels": labels}
+            return data
+
+        return data
+
+    @pytest.fixture
+    def iris_data(self):
+        iris_dataset = load_iris()
+        features, labels = iris_dataset.data, iris_dataset.target
+        K = len(iris_dataset.target_names)
+        N = features.shape[0]
+        pred_probs = self.generate_pred_probs(N, K, labels, noisy=True)
+        data = {"features": features, "pred_probs": pred_probs, "labels": labels}
+        return data
+
+    @pytest.fixture
+    def issue_manager(self, lab, make_data, monkeypatch):
+        data = make_data()
+        monkeypatch.setattr(lab._labels, "labels", data["labels"])
+        clustering_kwargs = {"eps": 2}
+        return UnderperformingGroupIssueManager(
+            datalab=lab, threshold=0.2, clustering_kwargs=clustering_kwargs
+        )
+
+    def test_find_issues_no_underperforming_group(self, issue_manager, make_data):
+        data = make_data()
+        features, labels, pred_probs = data["features"], data["labels"], data["pred_probs"]
+        N = len(labels)
+        issue_manager.find_issues(features=features, pred_probs=pred_probs)
+        issues, summary = issue_manager.issues, issue_manager.summary
+        assert np.sum(issues["is_underperforming_group_issue"]) == 0
+        expected_issue_mask = np.full(N, False, bool)
+        assert np.all(
+            issues["is_underperforming_group_issue"] == expected_issue_mask
+        ), "Issue mask should be correct"
+        expected_scores = np.ones(N)
+        np.testing.assert_allclose(
+            issues["underperforming_group_score"],
+            expected_scores,
+            err_msg="Scores should be correct",
+        )
+        assert summary["issue_type"][0] == "underperforming_group"
+        assert summary["score"][0] == pytest.approx(1.0, abs=0.01)
+        # Check with cluster_ids param
+        issue_manager.find_issues(features=features, pred_probs=pred_probs, cluster_ids=labels)
+        issues_with_clabels, summary_with_clabels = issue_manager.issues, issue_manager.summary
+        pd.testing.assert_frame_equal(issues_with_clabels, issues)
+        pd.testing.assert_frame_equal(summary_with_clabels, summary)
+
+    def test_find_issues(self, issue_manager, make_data):
+        RELATIVE_TOLERANCE = 1e-3
+        data = make_data(noisy=True)
+        features, labels, pred_probs = data["features"], data["labels"], data["pred_probs"]
+        N = len(labels)
+        issue_manager.find_issues(features=features, pred_probs=pred_probs)
+        issues, summary = issue_manager.issues, issue_manager.summary
+        assert np.sum(issues["is_underperforming_group_issue"]) == 100
+        expected_issue_mask = np.zeros(N, bool)
+        expected_issue_mask[labels == 0] = True
+        assert np.all(
+            issues["is_underperforming_group_issue"] == expected_issue_mask
+        ), "Issue mask should be correct"
+        expected_loss_ratio = 0.1429
+        expected_scores = np.ones(N)
+        expected_scores[labels == 0] = expected_loss_ratio
+        np.testing.assert_allclose(
+            issues["underperforming_group_score"],
+            expected_scores,
+            err_msg="Scores should be correct",
+            rtol=1e-3,
+        )
+        assert summary["issue_type"][0] == "underperforming_group"
+        assert summary["score"][0] == pytest.approx(expected_loss_ratio, rel=RELATIVE_TOLERANCE)
+        # Check with cluster_ids param
+        issue_manager.find_issues(features=features, pred_probs=pred_probs, cluster_ids=labels)
+        issues_with_clabels, summary_with_clabels = issue_manager.issues, issue_manager.summary
+        pd.testing.assert_frame_equal(issues_with_clabels, issues, rtol=RELATIVE_TOLERANCE)
+        pd.testing.assert_frame_equal(summary_with_clabels, summary, rtol=RELATIVE_TOLERANCE)
+        # With shifted cluster_ids
+        issue_manager.find_issues(features=features, pred_probs=pred_probs, cluster_ids=labels + 10)
+        issues_with_clabels, summary_with_clabels = issue_manager.issues, issue_manager.summary
+        pd.testing.assert_frame_equal(issues_with_clabels, issues, rtol=RELATIVE_TOLERANCE)
+        pd.testing.assert_frame_equal(summary_with_clabels, summary, rtol=RELATIVE_TOLERANCE)
+
+    def test_collect_info(self, issue_manager, make_data):
+        """Test some values in the info dict.
+
+        Mainly focused on the clustering info.
+        """
+        UNDERPERFORMING_CLUSTER_ID = 0
+        data = make_data(noisy=True)
+        features, pred_probs, labels = data["features"], data["pred_probs"], data["labels"]
+        issue_manager.find_issues(features=features, pred_probs=pred_probs)
+        info = issue_manager.info
+        assert "weighted_knn_graph" in info["statistics"]
+        assert "clustering" in info
+        # Check clustering info
+        clustering_info = info["clustering"]
+        assert clustering_info["algorithm"] == "DBSCAN"
+        assert clustering_info["params"]["metric"] == "precomputed"
+        assert clustering_info["stats"]["n_clusters"] == 4
+        # Test collect_info() with cluster_ids
+        issue_manager.find_issues(features=features, pred_probs=pred_probs, cluster_ids=labels)
+        info = issue_manager.info
+        assert "nearest_neighbor" not in info
+        assert "distance_to_nearest_neighbor" not in info
+        assert info["statistics"] == {}
+        # Check clustering info
+        clustering_info = info["clustering"]
+        assert clustering_info["algorithm"] is None
+        assert clustering_info["params"] == {}
+        assert clustering_info["stats"]["underperforming_cluster_id"] == UNDERPERFORMING_CLUSTER_ID
+        cluster_labels = clustering_info["stats"]["cluster_ids"]
+        issues = issue_manager.issues
+        issue_indices = issues.index[issues["is_underperforming_group_issue"]].values
+        assert np.all(
+            cluster_labels[issue_indices] == UNDERPERFORMING_CLUSTER_ID
+        ), "All samples with issue should belong to underperforming cluster"
+
+        np.testing.assert_equal(clustering_info["stats"]["cluster_ids"], labels)
+
+    def test_no_meaningful_clusters(self, issue_manager, make_data, monkeypatch):
+        np.random.seed(SEED)
+        data = make_data()
+        k = 10
+        N = 20
+        # Generate sparse data that cannot be clustered by DBSCAN
+        features = np.random.uniform(-100, 100, (N, 2))
+        # Dummy pred_probs and labels for running issue manager
+        pred_probs = data["pred_probs"][:N]
+        monkeypatch.setattr(issue_manager.datalab._labels, "labels", data["labels"][:N])
+        monkeypatch.setattr(issue_manager, "k", k)  # Ensure that k is smaller than N
+        exception_pattern = "No meaningful clusters"
+        with pytest.raises(ValueError, match=exception_pattern):
+            issue_manager.find_issues(features=features, pred_probs=pred_probs)
+        # Cluster labels passed containing all outliers
+        cluster_ids = np.full(N, -1)  # -1 is the outlier label for DBSCAN
+        with pytest.raises(ValueError, match=exception_pattern):
+            issue_manager.find_issues(
+                features=features, pred_probs=pred_probs, cluster_ids=cluster_ids
+            )
+        # Empty cluster ids
+        with pytest.raises(ValueError, match=exception_pattern):
+            issue_manager.find_issues(
+                features=features, pred_probs=pred_probs, cluster_ids=np.array([], dtype=int)
+            )
+
+    def test_min_cluster_samples(self, lab, issue_manager, make_data):
+        data = make_data()
+        features, pred_probs, labels = data["features"], data["pred_probs"], data["labels"]
+        labels[:3] = max(labels) + 1  # New cluster with very few samples
+        n_clusters = len(set(labels))
+        # Check if small cluster is filtered
+        issue_manager.find_issues(features=features, pred_probs=pred_probs, cluster_ids=labels)
+        clustering_info = issue_manager.info["clustering"]
+        assert clustering_info["stats"]["n_clusters"] == n_clusters - 1
+
+        # New issue manager to consider small cluster as well
+        issue_manager = UnderperformingGroupIssueManager(
+            datalab=lab, threshold=0.2, min_cluster_samples=3
+        )
+        issue_manager.find_issues(features=features, pred_probs=pred_probs, cluster_ids=labels)
+        clustering_info = issue_manager.info["clustering"]
+        assert clustering_info["stats"]["n_clusters"] == n_clusters
+
+    def test_find_issues_feature_subset(self, issue_manager, iris_data, monkeypatch):
+        features, pred_probs, labels = (
+            iris_data["features"],
+            iris_data["pred_probs"],
+            iris_data["labels"],
+        )
+        # TODO: Better asserts required. Ideally -> 3 clusters, 50 samples with issue.
+        monkeypatch.setattr(issue_manager.datalab._labels, "labels", labels)
+        monkeypatch.setattr(issue_manager, "clustering_kwargs", {"eps": 0.5})
+        # Find underperforming group based on one feature
+        single_feature = features[:, 0].reshape(-1, 1)
+        issue_manager.find_issues(features=single_feature, pred_probs=pred_probs)
+        assert np.sum(issue_manager.issues["is_underperforming_group_issue"]) == 16
+        # Find underperforming group based on two features
+        monkeypatch.setattr(issue_manager, "info", {})
+        issue_manager.find_issues(features=features[:, [1, 3]], pred_probs=pred_probs)
+        assert np.sum(issue_manager.issues["is_underperforming_group_issue"]) == 49
+        # Find underperforming group based on all features
+        monkeypatch.setattr(issue_manager, "info", {})
+        issue_manager.find_issues(features=features, pred_probs=pred_probs)
+        assert np.sum(issue_manager.issues["is_underperforming_group_issue"]) == 48
+
+    def test_knn_graph_change(self, issue_manager):
+        dist_matrix = np.random.randint(1, 5, size=(10, 10))
+        np.fill_diagonal(dist_matrix, 0)  # Make diagonal 0 to mimic distance matrix
+        knn_graph = sp.csr_matrix(dist_matrix)
+        nnz_before_clustering = knn_graph.nnz
+        issue_manager.perform_clustering(knn_graph)
+        nnz_after_clustering = knn_graph.nnz
+        assert nnz_before_clustering == nnz_after_clustering
+
+    def test_report(self, issue_manager, make_data):
+        data = make_data()
+        features, pred_probs = data["features"], data["pred_probs"]
+        issue_manager.find_issues(features=features, pred_probs=pred_probs)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+        )
+        assert isinstance(report, str)
+        assert (
+            "--------------- underperforming_group issues ---------------\n\nNumber of examples with this issue"
+        ) in report

tests/datalab/test_cleanvision_integration.py

@@ -0,0 +1,308 @@
+import matplotlib
+
+# Set non-interactive backend before importing pyplot to avoid GUI dependencies in CI
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+import pandas as pd
+
+from cleanlab import Datalab
+import cleanlab.datalab.internal.adapter.imagelab as imagelab
+
+LABEL_NAME = "label"
+IMAGE_NAME = "image"
+IMAGELAB_ISSUE_TYPES = [
+    "dark",
+    "light",
+    "low_information",
+    "odd_aspect_ratio",
+    "odd_size",
+    "grayscale",
+    "blurry",
+]
+SEED = 42
+
+
+class TestCleanvisionIntegration:
+    @pytest.fixture
+    def features(self, image_dataset):
+        np.random.seed(SEED)
+        return np.random.rand(len(image_dataset), 5)
+
+    @pytest.fixture
+    def num_imagelab_issues(self):
+        return 7
+
+    @pytest.fixture
+    def num_datalab_issues(self):
+        return 6
+
+    @pytest.fixture
+    def pred_probs(self, image_dataset):
+        np.random.seed(SEED)
+        return np.random.rand(len(image_dataset), 2)
+
+    @pytest.fixture
+    def set_plt_show(self, monkeypatch):
+        monkeypatch.setattr(plt, "show", lambda: None)
+
+    @pytest.mark.usefixtures("set_plt_show")
+    def test_imagelab_issues_checked(
+        self, image_dataset, pred_probs, features, capsys, num_imagelab_issues, num_datalab_issues
+    ):
+        datalab = Datalab(data=image_dataset, label_name=LABEL_NAME, image_key=IMAGE_NAME)
+        datalab.find_issues(pred_probs=pred_probs, features=features)
+        captured = capsys.readouterr()
+        assert (
+            "Finding dark, light, low_information, odd_aspect_ratio, odd_size, grayscale, blurry images"
+            in captured.out
+        )
+        # unable to check for non iid as feature space is too small, skipping it in interest of time
+        assert "Failed to check for these issue types: [NonIIDIssueManager]" in captured.out
+        assert len(datalab.issues) == len(image_dataset)
+
+        # add up imagelab + datalab issues
+        assert len(datalab.issues.columns) == (num_imagelab_issues + num_datalab_issues) * 2
+        assert len(datalab.issue_summary) == num_imagelab_issues + num_datalab_issues
+
+        all_keys = IMAGELAB_ISSUE_TYPES + [
+            "statistics",
+            "label",
+            "outlier",
+            "near_duplicate",
+            "class_imbalance",
+            "null",
+            "underperforming_group",
+            # "non_iid",
+            # Spurious correlations issue type is checked by default on image datasets
+            "spurious_correlations",
+        ]
+
+        assert set(all_keys) == set(datalab.info.keys())
+        datalab.report(show_all_issues=True)
+        captured = capsys.readouterr()
+
+        for issue_type in IMAGELAB_ISSUE_TYPES:
+            assert issue_type in captured.out
+
+        df = pd.DataFrame(
+            {
+                "issue_type": [
+                    "dark",
+                    "light",
+                    "low_information",
+                    "odd_aspect_ratio",
+                    "odd_size",
+                    "grayscale",
+                    "blurry",
+                    "label",
+                    "outlier",
+                    "near_duplicate",
+                    "class_imbalance",
+                    "null",
+                    "underperforming_group",
+                ],
+                "num_issues": [1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
+            }
+        )
+        expected_count = df.sort_values(by="issue_type")["num_issues"].tolist()
+        count = datalab.issue_summary.sort_values(by="issue_type")["num_issues"].tolist()
+        assert set(datalab.issue_summary["issue_type"].tolist()) == set(df["issue_type"].tolist())
+        assert count == expected_count
+        assert datalab.issue_summary["num_issues"].sum() == df["num_issues"].sum()
+
+    @pytest.mark.usefixtures("set_plt_show")
+    def test_imagelab_max_prevalence(
+        self,
+        image_dataset,
+        pred_probs,
+        features,
+        capsys,
+        num_datalab_issues,
+        monkeypatch,
+    ):
+        max_prevalence = 0
+        monkeypatch.setattr(imagelab, "IMAGELAB_ISSUES_MAX_PREVALENCE", max_prevalence)
+        datalab = Datalab(data=image_dataset, label_name=LABEL_NAME, image_key=IMAGE_NAME)
+
+        datalab.find_issues(pred_probs=pred_probs, features=features)
+        captured = capsys.readouterr()
+        assert (
+            "Finding dark, light, low_information, odd_aspect_ratio, odd_size, grayscale, blurry images"
+            in captured.out
+        )
+        assert (
+            f"from potential issues in the dataset as it exceeds max_prevalence={max_prevalence}"
+            in captured.out
+        )
+
+        issue_summary = datalab.get_issue_summary()
+        assert (
+            len(issue_summary) == 1 + num_datalab_issues
+        )  # adding 1 as no low_information issues present
+
+    def test_imagelab_issues_not_checked(
+        self, image_dataset, pred_probs, features, capsys, num_datalab_issues
+    ):
+        datalab = Datalab(data=image_dataset, label_name=LABEL_NAME)
+        datalab.find_issues(pred_probs=pred_probs, features=features)
+        captured = capsys.readouterr()
+        assert (
+            "Finding dark, light, low_information, odd_aspect_ratio, odd_size, grayscale, blurry images"
+            not in captured.out
+        )
+        assert len(datalab.issues) == len(image_dataset)
+        assert len(datalab.issues.columns) == num_datalab_issues * 2
+        assert len(datalab.issue_summary) == num_datalab_issues
+
+        all_keys = [
+            "statistics",
+            "label",
+            "outlier",
+            "near_duplicate",
+            "class_imbalance",
+            "null",
+            "underperforming_group",
+        ]
+
+        assert set(all_keys) == set(datalab.info.keys())
+        datalab.report(show_all_issues=True)
+        captured = capsys.readouterr()
+
+        for issue_type in IMAGELAB_ISSUE_TYPES:
+            assert issue_type not in captured.out
+
+    @pytest.mark.usefixtures("set_plt_show")
+    def test_incremental_issue_check(self, image_dataset, pred_probs, features, capsys):
+        datalab = Datalab(data=image_dataset, label_name=LABEL_NAME, image_key=IMAGE_NAME)
+        datalab.find_issues(pred_probs=pred_probs, features=features, issue_types={"label": {}})
+
+        assert len(datalab.issues) == len(image_dataset)
+        assert len(datalab.issues.columns) == 2
+        assert len(datalab.issue_summary) == 1
+
+        all_keys = ["statistics", "label"]
+
+        assert set(all_keys) == set(datalab.info.keys())
+
+        datalab.report(show_all_issues=True)
+        captured = capsys.readouterr()
+        assert "label" in captured.out
+
+        datalab.find_issues(issue_types={"image_issue_types": {"dark": {}}})
+
+        assert len(datalab.issues) == len(image_dataset)
+        assert len(datalab.issues.columns) == 4
+        assert len(datalab.issue_summary) == 2
+
+        all_keys = ["statistics", "label", "dark"]
+
+        assert set(all_keys) == set(datalab.info.keys())
+
+        datalab.report(show_all_issues=True)
+        captured = capsys.readouterr()
+        assert "label" in captured.out
+        assert "dark" in captured.out
+
+        with pytest.warns() as record:
+            datalab.find_issues(
+                issue_types={"image_issue_types": {"dark": {"threshold": 0.5}, "light": {}}}
+            )
+            assert len(record) == 3
+            assert (
+                "Overwriting columns ['is_dark_issue', 'dark_score'] in self.issues with columns from imagelab."
+                == record[0].message.args[0]
+            )
+            assert (
+                "Overwriting ['dark'] rows in self.issue_summary from imagelab."
+                == record[1].message.args[0]
+            )
+            assert "Overwriting key dark in self.info" == record[2].message.args[0]
+
+        assert len(datalab.issues) == len(image_dataset)
+        assert len(datalab.issues.columns) == 6
+        assert len(datalab.issue_summary) == 3
+
+        all_keys = ["statistics", "label", "dark", "light"]
+
+        assert set(all_keys) == set(datalab.info.keys())
+
+        datalab.report(show_all_issues=True)
+        captured = capsys.readouterr()
+        assert "label" in captured.out
+        assert "dark" in captured.out
+
+    @pytest.mark.usefixtures("set_plt_show")
+    def test_labels_not_required_for_imagelab_issues(
+        self, image_dataset, features, capsys, num_imagelab_issues
+    ):
+        datalab = Datalab(data=image_dataset, image_key=IMAGE_NAME)
+        datalab.find_issues()
+        captured = capsys.readouterr()
+        assert (
+            "Finding dark, light, low_information, odd_aspect_ratio, odd_size, grayscale, blurry images"
+            in captured.out
+        )
+        assert len(datalab.issues) == len(image_dataset)
+        assert len(datalab.issues.columns) == num_imagelab_issues * 2
+        assert len(datalab.issue_summary) == num_imagelab_issues
+
+        all_keys = IMAGELAB_ISSUE_TYPES + ["statistics"]
+
+        assert set(all_keys) == set(datalab.info.keys())
+        datalab.report(show_all_issues=True)
+        captured = capsys.readouterr()
+
+        for issue_type in IMAGELAB_ISSUE_TYPES:
+            assert issue_type in captured.out
+
+    @pytest.fixture
+    def lab(self, image_dataset):
+        lab = Datalab(data=image_dataset, label_name=LABEL_NAME, image_key=IMAGE_NAME)
+        lab.find_issues()
+        return lab
+
+    def test_get_summary(self, lab):
+        summary = lab.get_issue_summary("dark")
+        assert len(summary) == 1
+        num_issues = summary["num_issues"].values[0]
+        assert num_issues == 1
+
+    @pytest.mark.parametrize(
+        "list_method", ["list_possible_issue_types", "list_default_issue_types"]
+    )
+    def test_list_issue_type_method(self, image_dataset, lab, list_method):
+        method = getattr(lab, list_method)
+        issue_types = method()
+
+        # Check that Datalab without Imagelab injected has just a subset of possible/default issue types
+        minimal_lab = Datalab(data=image_dataset)
+        minimal_method = getattr(minimal_lab, list_method)
+        datalab_issue_types = minimal_method()
+        assert set(datalab_issue_types).issubset(set(issue_types))
+
+        # The additional issue types found by method should be the same as IMAGELAB_ISSUE_TYPES
+        assert set(issue_types).difference(datalab_issue_types) == set(IMAGELAB_ISSUE_TYPES)
+
+    @pytest.mark.issue1027
+    def test_get_issues(self, lab):
+        """
+        Test the `get_issues` method of the `lab` object.
+
+        This method checks if the columns returned by the `get_issues` method
+        match the expected columns for each issue type defined in `IMAGELAB_ISSUE_TYPES`.
+
+        Raises:
+            AssertionError: If the columns returned by `get_issues` do not match the expected columns.
+
+        """
+        test_condition = lambda s: set(lab.get_issues(s).columns) == set(
+            [f"{s}_score", f"is_{s}_issue"]
+        )
+        failed_assertions = [
+            issue_type for issue_type in IMAGELAB_ISSUE_TYPES if not test_condition(issue_type)
+        ]
+        assert (
+            len(failed_assertions) == 0
+        ), f"Tests for `get_issues` with these `issue_types` failed: {failed_assertions}"

tests/datalab/test_data.py

@@ -0,0 +1,246 @@
+import tempfile
+from unittest.mock import patch
+import pytest
+from cleanlab.datalab.internal.data import Data, DataFormatError, DatasetLoadError
+from datasets import Dataset, ClassLabel
+import numpy as np
+import hypothesis.strategies as st
+from hypothesis import given, assume, settings, HealthCheck
+
+from cleanlab.datalab.internal.task import Task
+
+
+NUM_COLS = 2
+
+
+@st.composite
+def multiclass_dataset_strategy(draw):
+    # Define strategies
+    int_feature_strategy = st.integers(min_value=-10, max_value=10)
+    float_feature_strategy = st.floats(min_value=-10, max_value=10)
+    column_name_strategy = st.text(
+        alphabet=st.characters(blacklist_categories=["Cs", "Cc", "Cn"]), min_size=5, max_size=5
+    )
+    column_data_strategy = st.one_of(int_feature_strategy, float_feature_strategy)
+
+    # Draw values
+    col_names = draw(
+        st.lists(column_name_strategy, min_size=NUM_COLS, max_size=NUM_COLS + 1, unique=True)
+    )
+    label_name = draw(st.sampled_from(col_names))
+    data = {
+        name: draw(st.lists(column_data_strategy, min_size=5, max_size=5)) for name in col_names
+    }
+    dataset = Dataset.from_dict(data)
+    dataset = dataset.rename_column(label_name, "label")
+
+    # Make assertions about drawn values
+    assume(len(set(dataset["label"])) > 1)
+
+    return dataset
+
+
+@st.composite
+def multilabel_dataset_strategy(draw):
+    # Define strategies
+    min_dataset_size = 5
+    max_dataset_size = 5
+    int_feature_strategy = st.integers(min_value=-10, max_value=10)
+    float_feature_strategy = st.floats(min_value=-10, max_value=10)
+    # Ensure column names do not include problematic characters
+    column_name_strategy = st.text(
+        alphabet=st.characters(blacklist_characters="\x00", min_codepoint=32, max_codepoint=126),
+        min_size=5,
+        max_size=5,
+    )
+    column_data_strategy = st.one_of(int_feature_strategy, float_feature_strategy)
+
+    # Draw values
+    col_names = draw(
+        st.lists(column_name_strategy, min_size=NUM_COLS, max_size=NUM_COLS + 1, unique=True)
+    )
+    label_name = draw(st.sampled_from(col_names))
+    # Ensure labels do not include problematic characters
+    classes_strategy = st.lists(
+        st.text(
+            alphabet=st.characters(
+                # The null character is problematic for some string operations, e.g. key lookup
+                blacklist_characters="\x00",
+                min_codepoint=32,
+                max_codepoint=126,
+            ),
+            min_size=2,
+            max_size=3,
+        ),
+        min_size=2,
+        max_size=3,
+        unique=True,
+    )
+    classes = draw(classes_strategy)
+    labels_strategy = st.lists(
+        st.lists(st.sampled_from(classes), min_size=1, max_size=3, unique=True),
+        min_size=min_dataset_size,
+        max_size=max_dataset_size,
+    )
+    data = {
+        name: draw(
+            st.lists(column_data_strategy, min_size=min_dataset_size, max_size=max_dataset_size)
+        )
+        for name in col_names
+    }
+    data[label_name] = draw(labels_strategy)
+    dataset = Dataset.from_dict(data)
+    dataset = dataset.rename_column(label_name, "label")
+
+    # Make assertions about drawn values
+    assume(len(set(l for labels in dataset["label"] for l in labels)) > 1)
+
+    return dataset
+
+
+@st.composite
+def dataset_strategy(draw, task=Task.CLASSIFICATION):
+    if task == Task.CLASSIFICATION:
+        return draw(multiclass_dataset_strategy())
+    elif task == Task.MULTILABEL:
+        return draw(multilabel_dataset_strategy())
+    else:
+        raise ValueError(f"Unsupported task: {task}")
+
+
+class TestData:
+    @pytest.fixture
+    def dataset_and_label_name(self):
+        label_name = "labels"
+
+        dataset = Dataset.from_dict({"image": [1, 2, 3], label_name: [0, 1, 0]})
+        return dataset, label_name
+
+    @given(dataset=dataset_strategy())
+    @settings(max_examples=10, suppress_health_check=[HealthCheck.too_slow])
+    def test_init_data_properties(self, dataset):
+        data = Data(data=dataset, task=Task.CLASSIFICATION, label_name="label")
+        assert data._data == dataset
+
+        # All elements in the _labels attribute are integers in the range [0, num_classes - 1]
+        num_classes = len(set(data.labels.label_map))
+        all_labels_are_ints = np.issubdtype(data.labels.labels.dtype, np.integer)
+        assert all_labels_are_ints, f"{data.labels.labels} should be a list of integers"
+        assert all(0 <= label < num_classes for label in data.labels.labels)
+
+        assert all(isinstance(label, int) for label in data.labels.label_map.keys())
+
+    def test_init_data(self, dataset_and_label_name):
+        dataset, label_name = dataset_and_label_name
+        data = Data(data=dataset, task=Task.CLASSIFICATION, label_name=label_name)
+
+        label_feature = dataset.features[label_name]
+        if isinstance(label_feature, ClassLabel):
+            classes = label_feature.names
+        else:
+            classes = sorted(dataset.unique(label_name))
+        assert data.class_names == classes
+
+    def test_init_data_from_list_of_dicts(self):
+        dataset = [{"X": 0, "label": 0}, {"X": 1, "label": 1}, {"X": 2, "label": 1}]
+        data = Data(data=dataset, task=Task.CLASSIFICATION, label_name="label")
+        assert isinstance(data._data, Dataset)
+
+    def test_init_raises_format_error(self):
+        data = np.random.rand(10, 2)
+        with pytest.raises(DataFormatError) as excinfo:
+            Data(data=data, task=Task.CLASSIFICATION, label_name="label")
+
+        expected_error_substring = "Unsupported data type: <class 'numpy.ndarray'>\n"
+        assert expected_error_substring in str(excinfo.value)
+
+    def test_init_raises_load_error(self):
+        improperly_aligned_data = {
+            "X": [0, 1, 2],
+            "label": [0, 1],
+        }
+        with pytest.raises(DatasetLoadError) as excinfo:
+            Data(data=improperly_aligned_data, task=Task.CLASSIFICATION, label_name="label")
+
+        expected_error_substring = "Failed to load dataset from <class 'dict'>.\n"
+        assert expected_error_substring in str(excinfo.value)
+
+    def test_not_equal_to_copy_or_non_data(self):
+        dataset = {"X": [0, 1, 2], "label": [0, 1, 2]}
+        data = Data(data=dataset, task=Task.CLASSIFICATION, label_name="label")
+        data_copy = Data(data=dataset, task=Task.CLASSIFICATION, label_name="label")
+        assert data != data_copy
+        assert data != dataset
+
+    def test_load_dataset_from_string(self, monkeypatch):
+        # Test with non-existent file
+        with pytest.raises(DatasetLoadError):
+            Data._load_dataset_from_string("non_existent_file.txt")
+
+        # Test with invalid extension
+        with tempfile.NamedTemporaryFile(suffix=".invalid") as temp_file:
+            with pytest.raises(DatasetLoadError):
+                Data._load_dataset_from_string(temp_file.name)
+
+        # Test with invalid external dataset identifier
+        with patch("datasets.load_dataset") as mock_load_dataset:
+            mock_load_dataset.side_effect = ValueError("Invalid external dataset identifier")
+            with pytest.raises(DatasetLoadError) as excinfo:
+                Data._load_dataset_from_string("invalid_external_dataset_name")
+
+            expected_error_substring = "Failed to load dataset from <class 'str'>.\n"
+            assert expected_error_substring in str(excinfo.value)
+
+        # Test with valid .txt, .csv, and .json files
+        test_data = [
+            (".txt", "sample text", "from_text"),
+            (".csv", "column1,column2\nvalue1,value2", "from_csv"),
+            (".json", '{"key": "value"}', "from_json"),
+        ]
+
+        mock_dataset = Dataset.from_dict({"y": [1, 2, 3]})
+        for ext, content, loader_func in test_data:
+            with tempfile.NamedTemporaryFile(suffix=ext, mode="w+t") as temp_file:
+                temp_file.write(content)
+                temp_file.flush()
+
+                # Make sure the correct loader function is called
+                def fake_loader(file_name):
+                    assert file_name == temp_file.name
+                    return mock_dataset
+
+                with monkeypatch.context() as mp:
+                    mp.setattr(Dataset, loader_func, fake_loader)
+                    loaded_dataset = Data._load_dataset_from_string(temp_file.name)
+                    assert isinstance(loaded_dataset, Dataset)
+                    assert loaded_dataset == mock_dataset
+
+        # Test with an external dataset
+        def fake_load_dataset(data_string):
+            if data_string == "external_dataset":
+                return mock_dataset
+
+            raise Exception("Not the expected dataset string")
+
+        with monkeypatch.context() as mp:
+            mp.setattr("datasets.load_dataset", fake_load_dataset)
+            loaded_dataset = Data._load_dataset_from_string("external_dataset")
+            assert isinstance(loaded_dataset, Dataset)
+            assert loaded_dataset == mock_dataset
+
+            with pytest.raises(DatasetLoadError) as excinfo:
+                Data._load_dataset_from_string("non_external_dataset")
+
+            expected_error_substring = "Failed to load dataset from <class 'str'>.\n"
+
+    @given(dataset=dataset_strategy(task=Task.CLASSIFICATION))
+    def test_label_map_is_lexicographically_ordered(self, dataset):
+        data = Data(data=dataset, task=Task.CLASSIFICATION, label_name="label")
+        label_map = data.labels.label_map
+        assert list(label_map.values()) == sorted(label_map.values())
+
+    @given(dataset=dataset_strategy(task=Task.MULTILABEL))
+    def test_label_map_is_lexicographically_ordered_multilabel(self, dataset):
+        data = Data(data=dataset, task=Task.MULTILABEL, label_name="label")
+        label_map = data.labels.label_map
+        assert list(label_map.values()) == sorted(label_map.values())

tests/datalab/test_data_issues.py

@@ -0,0 +1,52 @@
+import pytest
+from cleanlab.datalab.internal.data import Data
+from cleanlab.datalab.internal.data_issues import DataIssues, _ClassificationInfoStrategy
+from cleanlab.datalab.internal.task import Task
+
+
+class TestDataIssues:
+    labels = ["B", "A", "B"]
+    label_name = "labels"
+    strategy = _ClassificationInfoStrategy
+
+    @pytest.fixture
+    def data_issues(self):
+        data = Data(
+            data={self.label_name: self.labels},
+            task=Task.CLASSIFICATION,
+            label_name=self.label_name,
+        )
+        data_issues = DataIssues(data=data, strategy=self.strategy)
+        yield data_issues
+
+    def test_data_issues_init(self, data_issues):
+        assert hasattr(data_issues, "issues")
+        assert hasattr(data_issues, "issue_summary")
+        assert hasattr(data_issues, "info")
+
+    def test_statistics(self, data_issues):
+        stats = data_issues.statistics
+
+        assert stats == data_issues.info["statistics"]
+        assert stats["num_examples"] == 3, f"Incorrect number of examples: {stats['num_examples']}"
+        assert stats["class_names"] == ["A", "B"], f"Incorrect class names: {stats['class_names']}"
+        assert stats["num_classes"] == 2, f"Incorrect number of classes: {stats['num_classes']}"
+        assert stats["multi_label"] is False
+        assert (
+            stats["health_score"] is None
+        ), f"Health score should initially be None, but is {stats['health_score']}"
+
+    def test_get_info(self, data_issues):
+        with pytest.raises(ValueError):
+            data_issues.get_info("nonexistent_issue")
+        assert data_issues.get_info("statistics") == data_issues.info["statistics"]
+
+    def test_get_info_label(self, data_issues):
+        data_issues.info["label"] = {"given_label": [0, 1, 1], "predicted_label": [1, 0, 1]}
+        info = data_issues.get_info("label")
+
+        label_format_error_message = (
+            "get_info('label') should return the given label formatted with the class names"
+        )
+        assert info.get("given_label").tolist() == ["A", "B", "B"], label_format_error_message
+        assert info.get("predicted_label").tolist() == self.labels, label_format_error_message

tests/datalab/test_factory.py

@@ -0,0 +1,43 @@
+import pytest
+
+from cleanlab.datalab.internal.issue_manager_factory import register, REGISTRY
+from cleanlab import Datalab
+from cleanlab.datalab.internal.issue_manager.issue_manager import IssueManager
+from cleanlab.datalab.internal.task import Task
+
+
+@pytest.fixture
+def registry():
+    return REGISTRY
+
+
+def test_list_possible_issue_types(registry):
+    lab = Datalab(data=[], label_name=None)
+    issue_types = lab.list_possible_issue_types()
+    assert isinstance(issue_types, list)
+    possible_issues = [
+        "outlier",
+        "near_duplicate",
+        "non_iid",
+        "label",
+        "class_imbalance",
+        "underperforming_group",
+        "data_valuation",
+        "null",
+    ]
+    assert set(issue_types) == set(possible_issues)
+
+    test_key = "test_for_list_possible_issue_types"
+
+    class TestIssueManager(IssueManager):
+        issue_name = test_key
+
+    TestIssueManager = register(TestIssueManager)
+
+    issue_types = lab.list_possible_issue_types()
+    assert set(issue_types) == set(
+        possible_issues + [test_key]
+    ), "New issue type should be added to the list"
+
+    # Clean up
+    del registry[Task.CLASSIFICATION][test_key]

tests/datalab/test_imagelab_adapter.py

@@ -0,0 +1,25 @@
+from cleanlab.datalab.internal.adapter.imagelab import create_imagelab
+from cleanlab.datalab.internal.adapter.constants import DEFAULT_CLEANVISION_ISSUES
+
+
+class TestImagelabAdapater:
+    def test_create_imagelab(self, image_dataset):
+        imagelab = create_imagelab(image_dataset, "image")
+        assert imagelab is not None
+        assert hasattr(imagelab, "issues")
+        assert hasattr(imagelab, "issue_summary")
+        assert hasattr(imagelab, "info")
+
+    def test_imagelab_default_issue_types(self):
+        default_issues = DEFAULT_CLEANVISION_ISSUES
+        assert set(default_issues) == set(
+            [
+                "dark",
+                "light",
+                "low_information",
+                "odd_aspect_ratio",
+                "odd_size",
+                "grayscale",
+                "blurry",
+            ]
+        )

tests/datalab/test_init.py

@@ -0,0 +1,16 @@
+import sys
+import importlib
+from unittest.mock import patch
+
+
+def test_datalab_unavailable():
+    with patch.dict(sys.modules, {"cleanlab.datalab.datalab": ImportError("Mocked ImportError")}):
+        # Reload the module to trigger the import statement
+        import cleanlab
+
+        importlib.reload(cleanlab)
+
+        assert cleanlab.Datalab.message == (
+            "Datalab is not available due to missing dependencies. "
+            "To install Datalab, run `pip install 'cleanlab[datalab]'`."
+        )

tests/datalab/test_issue_finder.py

@@ -0,0 +1,244 @@
+import numpy as np
+import pytest
+
+from cleanlab import Datalab
+from cleanlab.datalab.internal.issue_finder import IssueFinder
+from cleanlab.datalab.internal.task import Task
+
+
+class TestIssueFinder:
+    task = Task.CLASSIFICATION
+
+    @pytest.fixture
+    def lab(self):
+        N = 30
+        K = 2
+        y = np.random.randint(0, K, size=N)
+        lab = Datalab(data={"y": y}, label_name="y")
+        return lab
+
+    @pytest.fixture
+    def issue_finder(self, lab):
+        return IssueFinder(datalab=lab, task=self.task)
+
+    def test_init(self, issue_finder):
+        assert issue_finder.verbosity == 1
+
+    @pytest.mark.parametrize("key", ["pred_probs", "features", "knn_graph"])
+    def test_get_available_issue_types_no_kwargs(self, issue_finder, key):
+        expected_issue_types = {"class_imbalance": {}}
+        issue_types = issue_finder.get_available_issue_types(**{key: None})
+        assert (
+            issue_types == expected_issue_types
+        ), "Only class_imbalance issue type for classification requires no kwargs"
+
+    @pytest.mark.parametrize(
+        "issue_types",
+        [
+            {"label": {}},
+            {"label": {"some_arg": "some_value"}},
+            {"label": {"some_arg": "some_value"}, "outlier": {}},
+            {"label": {}, "outlier": {}, "some_issue_type": {"some_arg": "some_value"}},
+            {},
+        ],
+    )
+    def test_get_available_issue_types_with_issue_types(self, issue_finder, issue_types):
+        available_issue_types = issue_finder.get_available_issue_types(issue_types=issue_types)
+        assert (
+            available_issue_types == issue_types
+        ), f"Failed to get available issue types with issue_types={issue_types}"
+
+    @pytest.mark.parametrize(
+        "keys, should_contain_underperforming_group",
+        [
+            # Test cases where 'pred_probs' is not provided, should all give False
+            (["features"], False),
+            (["knn_graph"], False),
+            (["cluster_ids"], False),
+            (["features", "knn_graph"], False),
+            (["features", "cluster_ids"], False),
+            (["knn_graph", "cluster_ids"], False),
+            (["features", "knn_graph", "cluster_ids"], False),
+            # Test cases where 'pred_probs' is provided should all give True
+            (["pred_probs", "features"], True),
+            (["pred_probs", "knn_graph"], True),
+            (["pred_probs", "cluster_ids"], True),
+            (["pred_probs", "features", "knn_graph"], True),
+            (["pred_probs", "features", "cluster_ids"], True),
+            (["pred_probs", "knn_graph", "cluster_ids"], True),
+            (["pred_probs", "features", "knn_graph", "cluster_ids"], True),
+            # only if other required keys are provided
+            (["pred_probs"], False),
+        ],
+        ids=lambda v: (
+            f"keys={v} "
+            if isinstance(v, list)
+            else ("> available" if v is True else "> unavailable")
+        ),
+    )
+    # Some warnings about preferring cluster_ids over knn_graph, or knn_graph over features can be ignored
+    @pytest.mark.filterwarnings(r"ignore:.*will (likely )?prefer.*:UserWarning")
+    # No other warnings should be allowed
+    @pytest.mark.filterwarnings("error")
+    def test_underperforming_group_availability_issue_1065(
+        self, issue_finder, keys, should_contain_underperforming_group
+    ):
+        """
+        Tests the availability of the 'underperforming_group' issue type based on the presence of 'pred_probs' and other required keys in the supplied arguments.
+
+        This test addresses issue #1065, where the mapping that decides which issue types to run based on the supplied arguments is incorrect.
+        Specifically, the 'underperforming_group' check should only be executed if 'pred_probs' and another required key are included in the supplied arguments.
+        See: https://github.com/cleanlab/cleanlab/issues/1065.
+
+        Parameters
+        ----------
+        keys : list
+            A list of keys to be included in the kwargs.
+        should_contain_underperforming_group : bool
+            A flag indicating whether the 'underperforming_group' issue type should be present in the available issue types.
+
+        Scenarios
+        ---------
+        Various combinations of 'features', 'pred_probs', 'knn_graph', and 'cluster_ids' are tested.
+
+        Asserts
+        -------
+        Ensures 'underperforming_group' is in the available issue types if 'pred_probs' and another required key are provided.
+        Ensures 'underperforming_group' is not in the available issue types if the required conditions are not met.
+        """
+        mock_value = object()  # Mock value to simulate presence of the required keys
+        kwargs = {key: mock_value for key in keys}
+
+        available_issue_types = issue_finder.get_available_issue_types(**kwargs)
+        if should_contain_underperforming_group:
+            assert (
+                "underperforming_group" in available_issue_types
+            ), "underperforming_group should be available if 'pred_probs' and another required key are provided"
+        else:
+            assert (
+                "underperforming_group" not in available_issue_types
+            ), "underperforming_group should not be available if the required conditions are not met"
+
+    def test_get_available_issue_types(self, issue_finder):
+        expected_issue_types = {"class_imbalance": {}}
+        # Test with no kwargs, no issue type expected to be returned
+        for key in ["pred_probs", "features", "knn_graph"]:
+            issue_types = issue_finder.get_available_issue_types(**{key: None})
+            assert (
+                issue_types == expected_issue_types
+            ), "Only class_imbalance issue type for classification requires no kwargs"
+
+        # Test with only issue_types, input should be
+        issue_types_dicts = [
+            {"label": {}},
+            {"label": {"some_arg": "some_value"}},
+            {"label": {"some_arg": "some_value"}, "outlier": {}},
+            {"label": {}, "outlier": {}, "some_issue_type": {"some_arg": "some_value"}},
+            {},
+        ]
+        for issue_types in issue_types_dicts:
+            available_issue_types = issue_finder.get_available_issue_types(issue_types=issue_types)
+            fail_msg = f"Failed to get available issue types with issue_types={issue_types}"
+            assert available_issue_types == issue_types, fail_msg
+
+        ## Test availability of underperforming_group issue type
+        only_features_available = {"features": np.random.random((10, 2))}
+        available_issue_types = issue_finder.get_available_issue_types(**only_features_available)
+        fail_msg = "underperforming_group should not be available if 'pred_probs' is not provided"
+        assert "underperforming_group" not in available_issue_types, fail_msg
+        features_and_pred_probs_available = {
+            **only_features_available,
+            "pred_probs": np.random.random((10, 2)),
+        }
+        available_issue_types = issue_finder.get_available_issue_types(
+            **features_and_pred_probs_available
+        )
+        fail_msg = "underperforming_group should be available if 'pred_probs' is provided"
+        assert "underperforming_group" in available_issue_types, fail_msg
+
+    def test_find_issues(self, issue_finder, lab):
+        N = len(lab.data)
+        K = lab.get_info("statistics")["num_classes"]
+        X = np.random.rand(N, 2)
+        pred_probs = np.random.rand(N, K)
+        pred_probs = pred_probs / pred_probs.sum(axis=1, keepdims=True)
+
+        data_issues = lab.data_issues
+        assert data_issues.issues.empty
+
+        issue_finder.find_issues(
+            features=X,
+            pred_probs=pred_probs,
+        )
+
+        assert not data_issues.issues.empty
+
+    def test_validate_issue_types_dict(self, issue_finder, monkeypatch):
+        issue_types = {
+            "issue_type_1": {f"arg_{i}": f"value_{i}" for i in range(1, 3)},
+            "issue_type_2": {f"arg_{i}": f"value_{i}" for i in range(1, 4)},
+        }
+        defaults_dict = issue_types.copy()
+
+        issue_types["issue_type_2"][
+            "arg_2"
+        ] = "another_value_2"  # Should be in default, but not affect the test
+        issue_types["issue_type_2"][
+            "arg_4"
+        ] = "value_4"  # Additional arg not in defaults should be allowed (ignored)
+
+        with monkeypatch.context() as m:
+            m.setitem(issue_types, "issue_type_1", {})
+            with pytest.raises(ValueError) as e:
+                issue_finder._validate_issue_types_dict(issue_types, defaults_dict)
+            assert all([string in str(e.value) for string in ["issue_type_1", "arg_1", "arg_2"]])
+
+
+class TestRegressionIssueFinder:
+    task = "regression"
+
+    @pytest.fixture
+    def lab(self):
+        N = 30
+        K = 2
+        y = np.random.randint(0, K, size=N)
+        lab = Datalab(data={"y": y}, label_name="y", task=self.task)
+        return lab
+
+    @pytest.fixture
+    def issue_finder(self, lab):
+        return IssueFinder(datalab=lab, task=Task.from_str(self.task))
+
+    def test_get_available_issue_types(self, issue_finder):
+        expected_issue_types = {}
+
+        # Test with no kwargs
+        for key in ["pred_probs", "features", "knn_graph"]:
+            issue_types = issue_finder.get_available_issue_types(**{key: None})
+            assert (
+                issue_types == expected_issue_types
+            ), "No issue type for regression requires no kwargs"
+
+        # Test with issue_types:
+        issue_types_dicts = [
+            {"label": {}},
+            {"label": {"some_arg": "some_value"}},
+            {"label": {"some_arg": "some_value"}, "outlier": {}},
+            {},
+        ]
+        supported_issue_types = ["label"]
+        for issue_types in issue_types_dicts:
+            available_issue_types = issue_finder.get_available_issue_types(issue_types=issue_types)
+            fail_msg = f"Failed to get available issue types with issue_types={issue_types}"
+            assert available_issue_types == issue_types, fail_msg
+
+        # Test with all kwargs
+        kwargs = {k: k for k in ["pred_probs", "features", "knn_graph"]}
+        kwargs["issue_types"] = {"label": {}}
+        available_issue_types = issue_finder.get_available_issue_types(**kwargs)
+        assert available_issue_types == {
+            "label": {
+                "predictions": "pred_probs",  # Expect the ModelOutput.argument class variable to replace the key
+                "features": "features",
+            },
+        }

tests/datalab/test_issue_manager.py

@@ -0,0 +1,117 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from cleanlab.datalab.internal.issue_manager import IssueManager
+from cleanlab.datalab.internal.issue_manager_factory import (
+    REGISTRY,
+    register,
+)
+from cleanlab.datalab.internal.task import Task
+
+
+class TestCustomIssueManager:
+    @pytest.mark.parametrize(
+        "score",
+        [0, 0.5, 1],
+        ids=["zero", "positive_float", "one"],
+    )
+    def test_make_summary_with_score(self, custom_issue_manager, score):
+        summary = custom_issue_manager.make_summary(score=score)
+
+        expected_summary = pd.DataFrame(
+            {
+                "issue_type": [custom_issue_manager.issue_name],
+                "score": [score],
+            }
+        )
+        assert pd.testing.assert_frame_equal(summary, expected_summary) is None
+
+    @pytest.mark.parametrize(
+        "score",
+        [-0.3, 1.5, np.nan, np.inf, -np.inf],
+        ids=["negative_float", "greater_than_one", "nan", "inf", "negative_inf"],
+    )
+    def test_make_summary_invalid_score(self, custom_issue_manager, score):
+        with pytest.raises(ValueError):
+            custom_issue_manager.make_summary(score=score)
+
+
+def test_register_custom_issue_manager(monkeypatch):
+    import io
+    import sys
+    import copy
+
+    # Prepare a copy of the original registry for cleanup
+    original_registry = copy.deepcopy(REGISTRY)
+
+    assert "foo" not in REGISTRY
+
+    class Foo(IssueManager):
+        issue_name = "foo"
+
+        def find_issues(self):
+            pass
+
+    Foo = register(Foo)
+
+    assert REGISTRY[Task.CLASSIFICATION].get("foo") == Foo
+
+    # Reregistering should overwrite the existing class, put print a warning
+
+    monkeypatch.setattr("sys.stdout", io.StringIO())
+
+    class NewFoo(IssueManager):
+        issue_name = "foo"
+
+        def find_issues(self):
+            pass
+
+    NewFoo = register(NewFoo)
+
+    assert REGISTRY[Task.CLASSIFICATION].get("foo") == NewFoo
+    assert all(
+        [
+            text in sys.stdout.getvalue()
+            for text in [
+                "Warning: Overwriting existing issue manager foo with ",
+                "NewFoo",
+                " for task classification.",
+            ]
+        ]
+    ), "Should print a warning"
+
+    # Reregistering for task should overwrite the existing class, put print a warning
+    class NewerFoo(IssueManager):
+        issue_name = "label"
+
+        def find_issues(self):
+            pass
+
+    NewerFoo = register(NewerFoo, task="classification")
+
+    assert REGISTRY[Task.CLASSIFICATION].get("label") == NewerFoo
+    assert all(
+        [
+            text in sys.stdout.getvalue()
+            for text in [
+                "Warning: Overwriting existing issue manager label with ",
+                "NewerFoo",
+                " for task classification.",
+            ]
+        ]
+    ), "Should print a warning"
+
+    # Registering any issue manager for another task is permitted
+    class Bar(IssueManager):
+        issue_name = "bar"
+
+        def find_issues(self):
+            pass
+
+    Bar = register(Bar, task="regression")
+
+    assert REGISTRY[Task.REGRESSION].get("bar") == Bar
+
+    # Cleanup
+    REGISTRY.update(original_registry)

tests/datalab/test_report.py

@@ -0,0 +1,237 @@
+from unittest.mock import Mock, patch
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from cleanlab import Datalab
+from cleanlab.datalab.internal.report import Reporter
+from cleanlab.datalab.internal.task import Task
+
+
+class TestReporter:
+    @pytest.fixture
+    def lab(self):
+        N = 30
+        K = 2
+        X = np.random.rand(N, K)
+        y = np.random.randint(0, K, size=N)
+        pred_probs = np.random.rand(N, K)
+        lab = Datalab(data={"y": y}, label_name="y")
+        lab.find_issues(features=X, pred_probs=pred_probs)
+        return lab
+
+    @pytest.fixture
+    def data_issues(self, lab):
+        return lab.data_issues
+
+    @pytest.fixture
+    def reporter(self, data_issues):
+        return Reporter(data_issues=data_issues, task=Task.CLASSIFICATION)
+
+    def test_init(self, reporter, data_issues):
+        assert reporter.data_issues == data_issues
+        assert reporter.verbosity == 1
+        assert reporter.include_description == True
+        assert reporter.show_summary_score == False
+
+        another_reporter = Reporter(data_issues=data_issues, task=Task.CLASSIFICATION, verbosity=2)
+        assert another_reporter.verbosity == 2
+
+    def test_report(self, reporter):
+        """Test that the report method works. It just wraps the get_report method in a print
+        statement."""
+        mock_get_report = Mock()
+
+        with patch("builtins.print") as mock_print:  # type: ignore
+            with patch.object(reporter, "get_report", mock_get_report):
+                reporter.report(num_examples=3)
+            mock_get_report.assert_called_with(num_examples=3)
+        mock_print.assert_called_with(mock_get_report.return_value)
+
+    @pytest.mark.parametrize("include_description", [True, False])
+    def test_get_report(self, reporter, data_issues, include_description, monkeypatch):
+        """Test that the report method works. Assuming we have two issue managers, each should add
+        their section to the report."""
+
+        mock_issue_manager = Mock()
+        mock_issue_manager.issue_name = "foo"
+        mock_issue_manager.report.return_value = "foo report"
+
+        class MockIssueManagerFactory:
+            @staticmethod
+            def from_str(*args, **kwargs):
+                return mock_issue_manager
+
+        monkeypatch.setattr(
+            "cleanlab.datalab.internal.report._IssueManagerFactory", MockIssueManagerFactory
+        )
+        mock_issues = pd.DataFrame(
+            {
+                "is_foo_issue": [False, True, False, False, False],
+                "foo_score": [0.6, 0.2, 0.7, 0.7, 0.8],
+            }
+        )
+        monkeypatch.setattr(data_issues, "issues", mock_issues)
+
+        mock_issue_summary = pd.DataFrame(
+            {
+                "issue_type": ["foo"],
+                "score": [0.6],
+                "num_issues": [1],
+            }
+        )
+
+        mock_info = {"foo": {"bar": "baz"}}
+
+        monkeypatch.setattr(data_issues, "issue_summary", mock_issue_summary)
+
+        reporter = Reporter(
+            data_issues=data_issues,
+            task=Task.CLASSIFICATION,
+            verbosity=0,
+            include_description=include_description,
+        )
+        monkeypatch.setattr(data_issues, "issues", mock_issues, raising=False)
+        monkeypatch.setattr(data_issues, "info", mock_info, raising=False)
+
+        monkeypatch.setattr(
+            reporter, "_write_summary", lambda *args, **kwargs: "Here is a lab summary\n\n"
+        )
+        report = reporter.get_report(num_examples=3)
+        expected_report = "\n\n".join(["Here is a lab summary", "foo report"])
+        assert report == expected_report
+
+    @pytest.mark.parametrize(
+        "show_all_issues, expected_report",
+        [
+            (True, "Here is a lab summary\n\nfoo report\n\n\nbar report"),
+            (False, "Here is a lab summary\n\nfoo report"),
+        ],
+    )
+    def test_show_all_issues(
+        self, reporter, data_issues, monkeypatch, show_all_issues, expected_report
+    ):
+        """Test that the report method works. Assuming we have two issue managers, each should add
+        their section to the report."""
+
+        mock_issue_manager_foo = Mock()
+        mock_issue_manager_foo.issue_name = "foo"
+        mock_issue_manager_foo.report.return_value = "foo report"
+
+        mock_issue_manager_bar = Mock()
+        mock_issue_manager_bar.issue_name = "bar"
+        mock_issue_manager_bar.report.return_value = "bar report"
+
+        class MockIssueManagerFactory:
+            @staticmethod
+            def from_str(*args, **kwargs):
+                name = kwargs["issue_type"]
+                issue_managers = {
+                    "foo": mock_issue_manager_foo,
+                    "bar": mock_issue_manager_bar,
+                }
+                issue_manager = issue_managers.get(name)
+                if issue_manager is None:
+                    raise ValueError(f"Unknown issue manager name: {name}")
+                return issue_manager
+
+        monkeypatch.setattr(
+            "cleanlab.datalab.internal.report._IssueManagerFactory", MockIssueManagerFactory
+        )
+        mock_issues = pd.DataFrame(
+            {
+                "is_foo_issue": [False, True, False, False, False],
+                "foo_score": [0.6, 0.2, 0.7, 0.7, 0.8],
+                "is_bar_issue": [False, False, False, False, False],
+                "bar_score": [0.7, 0.9, 0.8, 0.8, 0.8],
+            }
+        )
+        monkeypatch.setattr(data_issues, "issues", mock_issues)
+
+        # "bar" issue may be omitted in report, unless show_all_issues is True
+        mock_issue_summary = pd.DataFrame(
+            {
+                "issue_type": ["foo", "bar"],
+                "score": [0.6, 0.8],
+                "num_issues": [1, 0],
+            }
+        )
+
+        mock_info = {
+            "foo": {"foobar": "baz"},
+            "bar": {"barfoo": "bazbar"},
+        }
+
+        monkeypatch.setattr(data_issues, "issue_summary", mock_issue_summary)
+
+        reporter = Reporter(
+            data_issues=data_issues,
+            task="classification",
+            verbosity=0,
+            include_description=False,
+            show_all_issues=show_all_issues,
+        )
+        monkeypatch.setattr(data_issues, "issues", mock_issues, raising=False)
+        monkeypatch.setattr(data_issues, "info", mock_info, raising=False)
+
+        monkeypatch.setattr(
+            reporter, "_write_summary", lambda *args, **kwargs: "Here is a lab summary\n\n"
+        )
+        report = reporter.get_report(num_examples=3)
+        assert report == expected_report
+
+    summary = pd.DataFrame(
+        {
+            "issue_type": ["foo", "bar"],
+            "score": [0.6, 0.8],
+            "num_issues": [1, 0],
+        }
+    )
+
+    expected_filtered_summary = pd.DataFrame(
+        {
+            "issue_type": ["foo"],
+            "score": [0.6],
+            "num_issues": [1],
+        }
+    )
+
+    def test_summary_with_score(self, reporter, data_issues, monkeypatch):
+        """Test that the _write_summary method returns the expected output when show_summary_score is True.
+
+        It should include the score column in the summary and a note about what the score means.
+        """
+        mock_statistics = {"num_examples": 100, "num_classes": 5}
+        monkeypatch.setattr(data_issues, "get_info", lambda *args, **kwargs: mock_statistics)
+
+        expected_output = (
+            "Dataset Information: num_examples: 100, num_classes: 5\n\n"
+            + "Here is a summary of various issues found in your data:\n\n"
+            + self.expected_filtered_summary.to_string(index=False)
+            + "\n\n"
+            + "(Note: A lower score indicates a more severe issue across all examples in the dataset.)\n\n"
+            + "Learn about each issue: https://docs.cleanlab.ai/stable/cleanlab/datalab/guide/issue_type_description.html\n"
+            + "See which examples in your dataset exhibit each issue via: `datalab.get_issues(<ISSUE_NAME>)`\n\n"
+            + "Data indices corresponding to top examples of each issue are shown below.\n\n\n"
+        )
+
+        reporter.show_summary_score = True
+        assert reporter._write_summary(self.summary) == expected_output
+
+    def test_summary_without_score(self, reporter, data_issues, monkeypatch):
+        mock_statistics = {"num_examples": 100, "num_classes": 5}
+        monkeypatch.setattr(data_issues, "get_info", lambda *args, **kwargs: mock_statistics)
+
+        expected_output = (
+            "Dataset Information: num_examples: 100, num_classes: 5\n\n"
+            + "Here is a summary of various issues found in your data:\n\n"
+            + self.expected_filtered_summary.drop(columns=["score"]).to_string(index=False)
+            + "\n\n"
+            + "Learn about each issue: https://docs.cleanlab.ai/stable/cleanlab/datalab/guide/issue_type_description.html\n"
+            + "See which examples in your dataset exhibit each issue via: `datalab.get_issues(<ISSUE_NAME>)`\n\n"
+            + "Data indices corresponding to top examples of each issue are shown below.\n\n\n"
+        )
+
+        reporter.show_summary_score = False
+        assert reporter._write_summary(self.summary) == expected_output

tests/internal/neighbor/test_metric.py

@@ -0,0 +1,31 @@
+import numpy as np
+import pytest
+
+from cleanlab.internal.neighbor.metric import decide_default_metric
+
+
+@pytest.mark.parametrize(
+    "N",
+    [2, 10, 50, 100, 101],
+)
+def test_decide_default_metric_for_2d_and_3d_features(N):
+    # 2D and 3D features should always use the euclidean metric, disregarding the different implementations.
+    for M in [2, 3]:
+        X = np.random.rand(N, M)
+        metric = decide_default_metric(X)
+        if hasattr(metric, "__name__"):
+            error_msg = "The metric should be the string 'euclidean' for N > 100."
+            assert N <= 100, error_msg
+            metric = getattr(metric, "__name__")
+        assert metric == "euclidean"
+
+
+@pytest.mark.parametrize(
+    "M",
+    [4, 5, 10, 50, 100],
+)
+def test_decide_default_metric_for_high_dimensional_features(M):
+    # High-dimensional features should always use the cosine metric.
+    X = np.random.rand(100, M)
+    metric = decide_default_metric(X)
+    assert metric == "cosine"

tests/internal/neighbor/test_neighbor.py

@@ -0,0 +1,461 @@
+from typing import cast
+
+from hypothesis import given, strategies as st
+from hypothesis.extra.numpy import arrays
+import pytest
+import numpy as np
+from sklearn.neighbors import NearestNeighbors
+from scipy.sparse import csr_matrix
+
+
+from cleanlab.internal.neighbor import features_to_knn
+from cleanlab.internal.neighbor.knn_graph import (
+    correct_knn_distances_and_indices,
+    correct_knn_graph,
+    construct_knn_graph_from_index,
+    create_knn_graph_and_index,
+)
+
+
+@pytest.mark.parametrize(
+    "N",
+    [2, 10, 100, 101],
+    ids=lambda x: f"N={x}",
+)
+@pytest.mark.parametrize(
+    "M",
+    [2, 3, 4, 5, 10, 50, 100],
+    ids=lambda x: f"M={x}",
+)
+def test_features_to_knn(N, M):
+
+    features = np.random.rand(N, M)
+    if N >= 100:
+        features[-10:] = features[-11]  # Make the last 11 entries all identical, as an edge-case.
+    knn = features_to_knn(features)
+
+    assert isinstance(knn, NearestNeighbors)
+    knn = cast(NearestNeighbors, knn)
+    assert knn.n_neighbors == min(10, N - 1)
+    if M > 3:
+        metric = knn.metric
+        assert metric == "cosine"
+    else:
+        metric = knn.metric
+        if N <= 100:
+            assert hasattr(metric, "__name__")
+            metric = metric.__name__
+        assert metric == "euclidean"
+
+    if N >= 100:
+        distances, indices = knn.kneighbors(n_neighbors=10)
+        # Assert that the last 10 rows are identical to the 11th last row.
+        assert np.allclose(features[-10:], features[-11])
+        np.testing.assert_allclose(distances[-11:], 0, atol=1e-15)
+        # All the indices belong to the same example, so the set of indices should be the same.
+        # No guarantees about the order of the indices, but each point is not considered its own neighbor.
+        np.testing.assert_allclose(np.unique(indices[-11:]), np.arange(start=N - 11, stop=N))
+
+    # The knn object should be fitted to the features.
+    # TODO: This is not a good test, but it's the best we can do without exposing the internal state of the knn object.
+    # Assert is_
+    assert knn._fit_X is features
+
+
+def test_knn_kwargs():
+    """Check that features_to_knn passes additional keyword arguments to the NearestNeighbors constructor correctly."""
+    N, M = 100, 10
+    features = np.random.rand(N, M)
+    V = features.var(axis=0)
+    knn = features_to_knn(
+        features,
+        n_neighbors=6,
+        metric="seuclidean",
+        metric_params={"V": V},
+    )
+
+    assert knn.n_neighbors == 6
+    assert knn.radius == 1.0
+    assert (alg := knn.algorithm) == "auto"
+    assert knn.leaf_size == 30
+    assert knn.metric == "seuclidean"
+    assert knn.metric_params == {"V": V}
+    assert knn.p == 2
+    assert knn._fit_X is features  # Not a public attribute, bad idea to rely on this attribute.
+
+    # Attributes estimated from fitted data
+    assert knn.n_features_in_ == M
+    assert knn.effective_metric_params_ == {"V": V}
+    assert knn.effective_metric_ == "seuclidean"
+    assert knn.n_samples_fit_ == N
+    assert (
+        knn._fit_method == "ball_tree" if alg == "auto" else alg
+    )  # Should be one of ["kd_tree", "ball_tree" and "brute"], set with "algorithm"
+
+
+@pytest.mark.parametrize("metric", ["cosine", "euclidean"])
+def test_construct_knn_graph_from_index(metric):
+    N, k = 100, 10
+    knn = NearestNeighbors(n_neighbors=k, metric=metric)
+    X = np.random.rand(N, 10)
+    knn.fit(X)
+    knn_graph = construct_knn_graph_from_index(knn)
+
+    assert knn_graph.shape == (N, N)
+    assert knn_graph.nnz == N * k
+    assert knn_graph.dtype == np.float64
+    assert np.all(knn_graph.data >= 0)
+    assert np.all(knn_graph.indices >= 0)
+    assert np.all(knn_graph.indices < 100)
+
+    distances = knn_graph.data.reshape(N, k)
+    indices = knn_graph.indices.reshape(N, k)
+
+    # Assert all rows in distances are sorted
+    assert np.all(np.diff(distances, axis=1) >= 0)
+
+
+class TestKNNCorrection:
+    def test_knn_graph_corrects_missing_duplicates(self):
+        """Test that the KNN graph correction identifies missing duplicates and places them correctly."""
+        X = np.array(
+            [
+                [0, 0],
+                [0, 0],
+                [0, 0],
+                [1, 1],
+            ]
+        )
+
+        # k = 2
+        retrieved_distances = np.array(
+            [
+                [0, np.sqrt(2)],
+                [0, 0],
+                [0, 0],
+                [np.sqrt(2)] * 2,
+            ]
+        )
+        retrieved_indices = np.array(
+            [
+                [1, 3],
+                [0, 2],
+                [0, 1],
+                [0, 1],
+            ]
+        )
+
+        # Most of the retrieved distances are correct, except for the first row that has two exact duplicates
+        expected_distances = np.copy(retrieved_distances)
+        expected_distances[0] = [0, 0]
+        expected_indices = np.copy(retrieved_indices)
+        expected_indices[0] = [1, 2]
+
+        # Simulate an properly ordered KNN graph, which missed an exact duplicate in row 0
+        knn_graph = csr_matrix(
+            (retrieved_distances.ravel(), retrieved_indices.ravel(), np.arange(0, 9, 2)),
+            shape=(4, 4),
+        )
+        expected_knn_graph = csr_matrix(
+            (expected_distances.ravel(), expected_indices.ravel(), np.arange(0, 9, 2)),
+            shape=(4, 4),
+        )
+
+        # Test that the distances and indices are corrected
+        corrected_distances, corrected_indices = correct_knn_distances_and_indices(
+            X, retrieved_distances, retrieved_indices
+        )
+        np.testing.assert_array_equal(corrected_distances, expected_distances)
+        np.testing.assert_array_equal(corrected_indices, expected_indices)
+
+        # Test that the knn graph can be corrected as well
+        corrected_knn_graph = correct_knn_graph(X, knn_graph)
+        np.testing.assert_array_equal(corrected_knn_graph.toarray(), expected_knn_graph.toarray())
+
+    def test_knn_graph_corrects_order_of_duplicates(self):
+        """Ensure that KNN correction prioritizes duplicates correctly even when initial indices are out of order."""
+        X = np.array(
+            [
+                [0, 0],
+                [0, 0],
+                [0, 0],
+                [1, 1],
+            ]
+        )
+
+        retrieved_distances = np.array(
+            [
+                [np.sqrt(2), 0],  # Should be [0, 0]
+                [0, 0],
+                [0, 0],
+                [np.sqrt(2)] * 2,
+            ]
+        )
+        retrieved_indices = np.array(
+            [
+                [3, 1],  # Should be [1, 2]
+                [0, 2],
+                [1, 0],  # Should be [0, 1]
+                [0, 1],
+            ]
+        )
+
+        expected_distances = np.copy(retrieved_distances)
+        expected_distances[0] = [0, 0]
+
+        expected_indices = np.copy(retrieved_indices)
+        expected_indices[0] = [1, 2]
+        expected_indices[2] = [0, 1]
+
+        # Simulate an IMPROPERLY ordered KNN graph
+        knn_graph = csr_matrix(
+            (retrieved_distances.ravel(), retrieved_indices.ravel(), np.arange(0, 9, 2)),
+            shape=(4, 4),
+        )
+        expected_knn_graph = csr_matrix(
+            (expected_distances.ravel(), expected_indices.ravel(), np.arange(0, 9, 2)),
+            shape=(4, 4),
+        )
+
+        # Test that the distances and indices are corrected
+        corrected_distances, corrected_indices = correct_knn_distances_and_indices(
+            X, retrieved_distances, retrieved_indices
+        )
+        np.testing.assert_array_equal(corrected_distances, expected_distances)
+        np.testing.assert_array_equal(corrected_indices, expected_indices)
+
+        # Test that the knn graph can be corrected as well
+        corrected_knn_graph = correct_knn_graph(X, knn_graph)
+        np.testing.assert_array_equal(corrected_knn_graph.toarray(), expected_knn_graph.toarray())
+
+
+def noisy_euclidean_distance(x, y):
+    """Calculate Euclidean distance and add bias if the distance is exactly zero (points are identical)."""
+    distance = np.linalg.norm(x - y)
+    if all(x == y):
+        distance += 2
+    return distance
+
+
+def test_create_knn_graph_correctness():
+    """
+    Test to verify that the KNN graph creation and index correction handles duplicate points
+    and correctly calculates distances using a modified Euclidean distance metric that adds a large
+    bias for data points that are identical.
+    """
+
+    # Define a set of points with duplicates
+    X = np.array(
+        [
+            [0, 0],
+            [0, 0],
+            [0, 0],
+            [1, 1],
+        ]
+    )
+
+    # Define the expected distances and indices for k=3
+    expected_distances = np.array(
+        [
+            [0, 0, np.sqrt(2)],
+            [0, 0, np.sqrt(2)],
+            [0, 0, np.sqrt(2)],
+            [np.sqrt(2), np.sqrt(2), np.sqrt(2)],
+        ]
+    )
+    expected_indices = np.array(
+        [
+            [1, 2, 3],
+            [0, 2, 3],
+            [0, 1, 3],
+            [0, 1, 2],
+        ]
+    )
+
+    ### TESTING graph WITH corrections
+    knn_graph_corrected, _ = create_knn_graph_and_index(
+        features=X, n_neighbors=3, metric=noisy_euclidean_distance, correct_exact_duplicates=True
+    )
+    distances_corrected, indices_corrected = knn_graph_corrected.data.reshape(
+        4, 3
+    ), knn_graph_corrected.indices.reshape(4, 3)
+
+    # Assert the corrected graph matches expected values
+    np.testing.assert_array_equal(distances_corrected, expected_distances)
+    np.testing.assert_array_equal(indices_corrected, expected_indices)
+
+    ### TESTING graph WITHOUT corrections
+    # With the noisy metric, the exact duplicates may be missed
+    knn_graph, _ = create_knn_graph_and_index(
+        features=X, n_neighbors=3, metric=noisy_euclidean_distance, correct_exact_duplicates=False
+    )
+    distances, indices = knn_graph.data.reshape(4, 3), knn_graph.indices.reshape(4, 3)
+
+    # Check that all distances in the last row of the *incorrect* graph are identical
+    np.testing.assert_array_equal(
+        distances[-1], [np.sqrt(2)] * 3
+    )  # Don't confuse this with expected_distances[-1]
+
+    # Verify that the first neighbor for the first three points in the incorrect graph is the last point
+    np.testing.assert_array_equal(indices[:3, 0], [3] * 3)
+    np.testing.assert_array_equal(distances[:3, 0], [np.sqrt(2)] * 3)
+
+
+@given(
+    # A collection of data points to search over<
+    X=arrays(
+        dtype=np.float64,
+        shape=st.tuples(
+            st.integers(min_value=6, max_value=10), st.integers(min_value=2, max_value=3)
+        ),
+        elements=st.floats(min_value=-10, max_value=10),
+    ),
+    # Here are the K nearest neighbors we want to find
+    k=st.integers(min_value=1, max_value=5),
+)
+def test_create_knn_graph_properties(X, k):
+    """
+    Property-based test to verify that the KNN graph creation handles varying input sizes and
+    checks that indices and distances are consistent within the graph.
+    """
+    knn_graph, _ = create_knn_graph_and_index(
+        features=X, n_neighbors=k, metric=noisy_euclidean_distance, correct_exact_duplicates=False
+    )
+    distances, indices = knn_graph.data.reshape(X.shape[0], k), knn_graph.indices.reshape(
+        X.shape[0], k
+    )
+
+    # Corrected version to handle exact duplicates
+    knn_graph_corrected, _ = create_knn_graph_and_index(
+        features=X, n_neighbors=k, metric=noisy_euclidean_distance, correct_exact_duplicates=True
+    )
+    distances_corrected, indices_corrected = knn_graph_corrected.data.reshape(
+        X.shape[0], k
+    ), knn_graph_corrected.indices.reshape(X.shape[0], k)
+
+    # Testing properties
+    # Ensure no self-references unless k > number of points minus one
+    for i in range(X.shape[0]):
+        assert i not in indices[i] or k > X.shape[0] - 1
+
+    # Ensure distances are non-negative
+    assert np.all(distances >= 0), "All distances should be non-negative"
+    # but the corrected distances may be smaller
+    assert np.all(distances_corrected <= distances)
+
+
+@given(
+    # This point will be duplicated
+    base_point=arrays(
+        dtype=np.float64,
+        shape=(2,),
+        elements=st.floats(min_value=-10, max_value=10, allow_subnormal=False),
+    ),
+    # This is how many instances there are of the duplicated point
+    num_duplicates=st.integers(min_value=2, max_value=5),
+    # Here are other points which aren't duplicates, will be post-processed to eliminate exact duplicates,
+    # so that they won't affect the duplicate results
+    extra_points=arrays(
+        dtype=np.float64,
+        shape=st.tuples(st.integers(min_value=11, max_value=20), st.just(2)),
+        elements=st.floats(min_value=15, max_value=20, allow_subnormal=False),
+        unique=True,
+    ),
+    # Here are the K nearest neighbors we want to find
+    k=st.integers(min_value=1, max_value=10),
+)
+def test_knn_graph_duplicate_handling(base_point, num_duplicates, extra_points, k):
+    """
+    Test to ensure that KNN graph handles duplicates properly by comparing graphs with and without
+    exact duplicate corrections.
+    """
+    # Before the test, ensure that the base_point is not a part of the extra_points (it's ok throw that point out)
+    if np.any(_dup := (extra_points == base_point).all(axis=1)):
+        extra_points = np.delete(extra_points, np.where(_dup)[0][0], axis=0)
+
+    # Create a dataset with duplicates of a single point and some extra distinct points
+    X = np.vstack([np.tile(base_point, (num_duplicates, 1)), extra_points])
+
+    # Run KNN without correcting duplicates
+    knn_graph, _ = create_knn_graph_and_index(
+        features=X, n_neighbors=k, metric=noisy_euclidean_distance, correct_exact_duplicates=False
+    )
+    distances, indices = knn_graph.data.reshape(X.shape[0], k), knn_graph.indices.reshape(
+        X.shape[0], k
+    )
+
+    # Run KNN with correcting duplicates
+    knn_graph_corrected, _ = create_knn_graph_and_index(
+        features=X, n_neighbors=k, metric=noisy_euclidean_distance, correct_exact_duplicates=True
+    )
+    distances_corrected, indices_corrected = knn_graph_corrected.data.reshape(
+        X.shape[0], k
+    ), knn_graph_corrected.indices.reshape(X.shape[0], k)
+
+    # Check two properties of the graphs, once corrected for duplicates
+    # 1. Check that duplicate points have the same neighbors in the corrected graph,
+    #    they should be their mutual closest neighbors
+
+    # To simplify comparisons across rows, include the row id (omitted in the knn_graphs)
+    duplicate_ids = np.arange(num_duplicates)
+    # Reshape the duplicate_ids array to a 2D array for later concatenation
+    reshaped_duplicate_ids = duplicate_ids.reshape(-1, 1)
+    # Get the nearest neighbors for each duplicate, excluding itself
+    nearest_neighbors = indices_corrected[:num_duplicates, : (num_duplicates - 1)]
+
+    # Concatenate the duplicate ids with their corresponding nearest neighbors
+    # This forms a 2D array where each row represents a duplicate and its nearest neighbors
+    # Note that in this test, the neighbors of interest are supposed to be duplicates themselves
+    points_and_neighbors = np.hstack((reshaped_duplicate_ids, nearest_neighbors))
+
+    # Define a function to calculate precision
+    # Precision is defined as the number of true positives divided by the number of true positives plus the number of false positives
+    def calculate_precision(test_set, actual_set):
+        true_positives = np.intersect1d(test_set, actual_set)
+        return len(true_positives) / len(test_set)
+
+    # All the points and their neighbors belong to the same set of duplicates.
+    for row in points_and_neighbors:
+        # Assert that the precision of the nearest neighbors with respect to the actual duplicates is 1
+        # This means that all nearest neighbors are actual duplicates
+        assert calculate_precision(row, duplicate_ids) == 1
+
+    # 2. Distances for duplicates are corrected (should be zeros)
+    # Without correcting for duplicates in this test, we assume some distances are non-zero
+    # We sum the distances for the first (duplicates-1) items
+    # If the distances were correctly calculated, the sum should be zero
+    # Therefore, if the sum is greater than zero, the distances were not correctly calculated
+    uncorrected_distances_sum = distances[:num_duplicates, : (num_duplicates - 1)].sum(axis=1)
+    assert any(
+        uncorrected_distances_sum > 0
+    ), "Uncorrected distances for duplicates are not greater than zero"
+    # With correction for duplicates, the distances should be zero
+    # We check this by comparing the corrected distances for the first (duplicates-1) items to zero
+    corrected_distances = distances_corrected[: (num_duplicates - 1), : (num_duplicates - 1)]
+    np.testing.assert_array_equal(
+        corrected_distances, 0, "Corrected distances for duplicates are not zero"
+    )
+
+
+def test_construct_knn_then_correct_knn_graph_does_the_same_work():
+    features = np.random.rand(1000, 2)
+    features[10:20] = features[10]  # Make the 10th to 20th rows identical
+    metric = noisy_euclidean_distance
+    n_neighbors = 50
+
+    # Construct the index and knn_graph separately, the correction should happen during the knn_graph construction
+    knn = features_to_knn(features, n_neighbors=n_neighbors, metric=metric)
+    knn_graph_from_index = construct_knn_graph_from_index(knn)  # Without correction
+    knn_graph_from_index_with_correction = construct_knn_graph_from_index(
+        knn, correction_features=features
+    )
+    knn_graph, _ = create_knn_graph_and_index(
+        features=features, n_neighbors=n_neighbors, metric=metric
+    )
+
+    # knn_graph has correction
+    np.testing.assert_array_equal(
+        knn_graph_from_index_with_correction.toarray(), knn_graph.toarray()
+    )
+    # knn_graph_from_index does not have correction
+    assert not np.all(knn_graph_from_index.toarray() == knn_graph.toarray())

tests/internal/test_numerics.py

@@ -0,0 +1,71 @@
+import pytest
+import numpy as np
+
+from cleanlab.internal.numerics import softmax
+
+
+class TestSoftmax:
+    def test_basic_softmax(self):
+        input_arr = np.array([1.0, 2.0, 3.0])
+        output = softmax(input_arr)
+        expected_output = np.array([0.09003057, 0.24472847, 0.66524096])
+        assert np.isclose(np.sum(output), 1.0)
+        assert np.allclose(output, expected_output)
+
+    def test_temperature_effect(self):
+        input_arr = np.array([1.0, 2.0, 3.0])
+        output_high_temp = softmax(input_arr, temperature=5.0)
+        output_low_temp = softmax(input_arr, temperature=0.1)
+
+        expected_high_temp = np.array([0.2693075, 0.32893292, 0.40175958])
+        expected_low_temp = np.array([2.06106005e-09, 4.53978686e-05, 9.99954600e-01])
+
+        assert np.allclose(output_high_temp, expected_high_temp)
+        assert np.allclose(output_low_temp, expected_low_temp)
+
+    def test_axis(self):
+        input_arr = np.array(
+            [
+                [1, 2, 3],  # unit step
+                [4, 5, 6],  # unit step
+                [7, 8, 10],  # non-unit step
+            ]
+        )
+        output = softmax(input_arr, axis=1)
+
+        expected_output = np.array(
+            [
+                [0.09003057, 0.24472847, 0.66524096],  # unit step
+                [0.09003057, 0.24472847, 0.66524096],  # unit step
+                [0.04201007, 0.1141952, 0.84379473],  # non-unit step
+            ]
+        )
+
+        assert np.allclose(output, expected_output)
+
+    @pytest.mark.parametrize(
+        "input_arr, expected_output",
+        [
+            (np.array([1.0, 2.0, 3.0]) + 1000, np.array([0.09003057, 0.24472847, 0.66524096])),
+            (np.array([1e3, 2e3, 3e3]), np.array([0, 0, 1])),
+        ],
+    )
+    def test_shift(self, input_arr, expected_output):
+        # Without shift, softmax overflows and gets a RuntimeWarning, but just returns nan
+        with pytest.warns(RuntimeWarning):
+            output_no_shift = softmax(input_arr, shift=False)
+        assert np.isnan(output_no_shift).all()
+
+        output_shift = softmax(input_arr, shift=True)
+        assert np.allclose(output_shift, expected_output)
+
+    @pytest.mark.parametrize(
+        "input_arr, expected_output",
+        [
+            (np.array([0, -np.inf, -np.inf]), np.array([1.0, 0.0, 0.0])),
+            (np.array([-np.inf, 0, 1]), np.array([0.0, 0.26894142, 0.73105858])),
+        ],
+    )
+    def test_special_values(self, input_arr, expected_output):
+        output = softmax(input_arr)
+        assert np.allclose(output, expected_output)

tests/spurious_correlation/test_correlation_visualizer.py

@@ -0,0 +1,48 @@
+from unittest.mock import patch
+
+from cleanlab.datalab.internal.adapter.imagelab import CorrelationVisualizer
+
+VIZMANAGER_IMPORT_PATH = "cleanvision.utils.viz_manager.VizManager"
+
+
+class TestCorrelationVisualizer:
+    def test_correlation_visualizer_init(self):
+        with patch(VIZMANAGER_IMPORT_PATH) as mock_viz_manager:
+            visualizer = CorrelationVisualizer()
+            assert visualizer.viz_manager == mock_viz_manager
+
+    def test_visualize(self):
+        with patch(VIZMANAGER_IMPORT_PATH) as mock_viz_manager:
+            visualizer = CorrelationVisualizer()
+
+            images = ["image1", "image2", "image3"]
+            title_info = {"scores": ["score1", "score2", "score3"]}
+            ncols = 2
+            cell_size = (2, 2)
+
+            visualizer.visualize(images, title_info)
+
+            mock_viz_manager.individual_images.assert_called_once_with(
+                images=images,
+                title_info=title_info,
+                ncols=ncols,
+                cell_size=cell_size,
+            )
+
+    def test_visualize_custom_params(self):
+        with patch(VIZMANAGER_IMPORT_PATH) as mock_viz_manager:
+            visualizer = CorrelationVisualizer()
+
+            images = ["image1", "image2", "image3"]
+            title_info = {"scores": ["score1", "score2", "score3"]}
+            ncols = 3
+            cell_size = (3, 3)
+
+            visualizer.visualize(images, title_info, ncols=ncols, cell_size=cell_size)
+
+            mock_viz_manager.individual_images.assert_called_once_with(
+                images=images,
+                title_info=title_info,
+                ncols=ncols,
+                cell_size=cell_size,
+            )

tests/spurious_correlation/test_spurious_correlation.py

@@ -0,0 +1,608 @@
+"""
+This script performs a spurious correlation test to detect and measure unintended associations
+between different transformations of images and their labels in a synthetic dataset.
+
+The process involves:
+1. Generating images with various shapes (circles and squares).
+2. Applying different filters (dark, blurry, and odd aspect ratio) to these images.
+3. Creating datasets with these filtered images and comparing them to a standard set of images without filters.
+
+The goal is to evaluate whether the application of specific filters results in lower correlation scores,
+indicating that the transformations introduce spurious correlations that might not be present in the original data.
+This helps ensure that these filters don't create misleading patterns that could affect the performance and reliability
+of machine learning models trained on such data.
+
+The test is implemented using the cleanlab library using Datalab module, which helps identify and quantify these spurious correlations.
+"""
+
+import re
+import numpy as np
+from PIL import Image, ImageDraw, ImageEnhance, ImageFilter
+import random
+from datasets import Dataset
+import pytest
+from cleanlab import Datalab
+import contextlib
+import io
+from unittest import mock
+from cleanlab.datalab.internal.adapter.constants import SPURIOUS_CORRELATION_ISSUE
+
+seed = 42
+np.random.seed(seed=seed)
+
+
+def create_base_image(size=(64, 64), background_color=(255, 255, 255)):
+    """
+    Creates a base image with the given size and background color.
+
+    Args:
+        size (tuple): The size of the image (width, height).
+        background_color (tuple): The background color of the image (RGB).
+
+    Returns:
+        Image: The created base image.
+    """
+    return Image.new("RGB", size, background_color)
+
+
+def draw_shape(draw, shape, color, offset_x, offset_y, shape_size):
+    """
+    Draws a specified shape on the image.
+
+    Args:
+        draw (ImageDraw): The drawing context.
+        shape (str): The shape to draw ('circle' or 'square').
+        color (tuple): The color of the shape (RGB).
+        offset_x (int): The x offset of the shape.
+        offset_y (int): The y offset of the shape.
+        shape_size (int): The size of the shape.
+    """
+    if shape == "circle":
+        draw.ellipse(
+            [(offset_x, offset_y), (offset_x + shape_size, offset_y + shape_size)], fill=color
+        )
+    elif shape == "square":
+        draw.rectangle(
+            [(offset_x, offset_y), (offset_x + shape_size, offset_y + shape_size)], fill=color
+        )
+
+
+def add_noise(image):
+    """
+    Adds random noise to the image.
+
+    Args:
+        image (Image): The image to add noise to.
+
+    Returns:
+        Image: The image with added noise.
+    """
+    np_img = np.array(image)
+    noise = np.random.normal(0, 0.5, np_img.shape).astype(np.uint8)
+    np_img = np.clip(np_img + noise, 0, 255)
+    return Image.fromarray(np_img)
+
+
+def create_image(shape, color, size=(64, 64), background_color=(255, 255, 255)):
+    """
+    Creates an image with a given shape and color.
+
+    Args:
+        shape (str): The shape to draw ('circle' or 'square').
+        color (tuple): The color of the shape (RGB).
+        size (tuple): The size of the image (width, height).
+        background_color (tuple): The background color of the image (RGB).
+
+    Returns:
+        Image: The generated image with the shape.
+    """
+    img = create_base_image(size, background_color)
+    draw = ImageDraw.Draw(img)
+
+    offset_x, offset_y, shape_size = randomize_shape_position_and_size(size)
+
+    draw_shape(draw, shape, color, offset_x, offset_y, shape_size)
+
+    img = add_noise(img)
+
+    return img
+
+
+def randomize_shape_position_and_size(size):
+    """
+    Randomizes the position and size of the shape.
+
+    Args:
+        size (tuple): The size of the image (width, height).
+
+    Returns:
+        tuple: The x offset, y offset, and size of the shape.
+    """
+    max_offset = 10
+    offset_x = random.randint(0, max_offset)
+    offset_y = random.randint(0, max_offset)
+    shape_size = random.randint(20, size[0] - 20)
+
+    return offset_x, offset_y, shape_size
+
+
+# Transformations
+def apply_dark(image):
+    """Decreases brightness of the image."""
+    enhancer = ImageEnhance.Brightness(image)
+    return enhancer.enhance(0.3)
+
+
+def apply_blurry(image, radius=20):
+    """Applies Gaussian blur to the image."""
+    return image.filter(ImageFilter.GaussianBlur(radius=radius))
+
+
+def apply_identity(image):
+    """Returns the unchanged image."""
+    return image
+
+
+def apply_odd_aspect_ratio(image):
+    """Changes the aspect ratio to make the image tall and skinny."""
+    return image.resize((32, 128))
+
+
+def generate_backgrounds(num_variations=5):
+    """
+    Generates a list of random background variations.
+
+    Args:
+        num_variations (int): The number of background variations to generate.
+
+    Returns:
+        list: A list of tuples with brightness and color.
+    """
+    backgrounds = []
+    for _ in range(num_variations):
+        brightness = random.uniform(0.2, 1.0)
+        color = tuple(random.choices(range(256), k=3))
+        backgrounds.append((brightness, color))
+    return backgrounds
+
+
+def apply_background(image, brightness, color):
+    """
+    Applies a background color to the image with specified brightness.
+
+    Args:
+        image (Image): The image to apply the background to.
+        brightness (float): The brightness level of the background.
+        color (tuple): The RGB color of the background.
+
+    Returns:
+        Image: The image with the background applied.
+    """
+    enhancer = ImageEnhance.Brightness(Image.new("RGB", image.size, color))
+    background = enhancer.enhance(brightness)
+    return Image.alpha_composite(background.convert("RGBA"), image.convert("RGBA")).convert("RGB")
+
+
+def apply_filter(image, filter_function):
+    """
+    Applies a random filter from the filter functions to the image.
+
+    Args:
+        image (Image): The image to apply the filter to.
+        filter_function (function): A filter function to apply.
+
+    Returns:
+        Image: The filtered image.
+    """
+    filtered_img = filter_function(image)
+    return filtered_img
+
+
+def generate_image_with_background(shape, color, filter_function, backgrounds):
+    """
+    Generates an image with a specified shape and color, applies a random filter,
+    and then applies a random background.
+
+    Args:
+        shape (str): The shape to draw.
+        color (tuple): The color of the shape.
+        filter_function (function): A filter function to apply.
+        backgrounds (list): A list of background variations.
+
+    Returns:
+        tuple: The generated image, filter type, brightness, and background color.
+    """
+    img = create_image(shape, color)
+    filtered_img = apply_filter(img, filter_function)
+    brightness, bg_color = random.choice(backgrounds)
+    background_img = apply_background(filtered_img, brightness, bg_color)
+    return background_img
+
+
+def get_filter_functions():
+    """
+    Returns a dictionary of available filter functions.
+
+    Returns:
+        dict: A dictionary mapping filter names to filter functions.
+    """
+    filter_functions_map = {
+        "dark": apply_dark,
+        "blurry": apply_blurry,
+        "identity": apply_identity,
+        "odd_aspect_ratio": apply_odd_aspect_ratio,
+    }
+    return filter_functions_map
+
+
+def get_filter_functions_without_identity():
+    """
+    Returns a dictionary of available filter functions, excluding 'identity'.
+
+    Returns:
+        dict: A dictionary mapping filter names to filter functions.
+    """
+    filter_functions_map = get_filter_functions()
+    del filter_functions_map["identity"]
+    return filter_functions_map
+
+
+def generate_dataset(
+    num_images_per_class=50,
+    num_background_variations=5,
+    circle_filter="identity",
+    square_filter="identity",
+):
+    """
+    Generates a toy dataset with images and corresponding labels.
+
+    Args:
+        num_images_per_class (int): The number of images per class.
+        num_background_variations (int): The number of background variations.
+        circle_filter (str): The filter to apply to circle images.
+        square_filter (str): The filter to apply to square images.
+
+    Returns:
+        Dataset: The generated dataset.
+    """
+    shapes = ["circle", "square"]
+    filter_functions = [circle_filter, square_filter]
+    colors = [
+        (255, 0, 0),  # Red
+        (0, 255, 0),  # Green
+        (0, 0, 255),  # Blue
+        (255, 255, 0),  # Yellow
+        (255, 0, 255),  # Magenta
+        (0, 255, 255),  # Cyan
+        (192, 192, 192),  # Gray
+    ]
+    filter_functions_map = get_filter_functions()
+    filter_functions = [filter_functions_map[circle_filter], filter_functions_map[square_filter]]
+    backgrounds = generate_backgrounds(num_background_variations)
+
+    data = []
+    labels = []
+
+    for shape, filter_function in zip(shapes, filter_functions):
+        for _ in range(num_images_per_class):
+            color = random.choice(colors)
+            background_img = generate_image_with_background(
+                shape, color, filter_function, backgrounds
+            )
+
+            data.append(background_img)
+            labels.append(shape)
+
+    dataset = Dataset.from_dict({"image": data, "label": labels})
+    return dataset
+
+
+def get_property_score(df, property):
+    """
+    Retrieves the score for a specific property from the dataframe.
+
+    Args:
+        df (DataFrame): The dataframe containing property scores.
+        property (str): The property to retrieve the score for.
+
+    Returns:
+        float: The score for the specified property.
+    """
+    return df.loc[df["property"] == property, "score"].iloc[0]
+
+
+def get_scores(df):
+    """
+    Retrieves scores for all relevant properties from the dataframe.
+
+    Args:
+        df (DataFrame): The dataframe containing property scores.
+
+    Returns:
+        dict: A dictionary with property names as keys and their scores as values.
+    """
+    filter_functions_map = get_filter_functions_without_identity()
+    properties_of_interest = [prop + "_score" for prop in filter_functions_map.keys()]
+    standard_label_uncorrelatedness_scores = {
+        prop: get_property_score(df, prop) for prop in properties_of_interest
+    }
+    return standard_label_uncorrelatedness_scores
+
+
+def get_label_uncorrelatedness_scores(circle_filter="identity", square_filter="identity"):
+    """
+    Generates a dataset and computes label uncorrelatedness scores.
+
+    Args:
+        circle_filter (str): The filter to apply to circle images.
+        square_filter (str): The filter to apply to square images.
+
+    Returns:
+        dict: A dictionary with property names as keys and their uncorrelatedness scores as values.
+    """
+    dataset = generate_dataset(circle_filter=circle_filter, square_filter=square_filter)
+    lab = Datalab(data=dataset, label_name="label", image_key="image")
+    lab.find_issues()
+    label_uncorrelatedness_scores = lab.get_info("spurious_correlations")["correlations_df"]
+    return get_scores(label_uncorrelatedness_scores)
+
+
+@pytest.mark.parametrize(
+    "test_attribute",
+    [
+        "dark",
+        "blurry",
+        "odd_aspect_ratio",
+    ],
+)
+def test_uncorrelatedness_scores_against_standard(test_attribute):
+    """
+    Tests that label uncorrelatedness scores for specific filters are lower than standard scores.
+
+    Asserts:
+        AssertionError: If any of the specific filter scores are not lower than the standard scores.
+    """
+    standard_label_uncorrelatedness_scores = get_label_uncorrelatedness_scores()
+    attribute_filter_label_uncorrelatedness_scores = get_label_uncorrelatedness_scores(
+        circle_filter=f"{test_attribute}"
+    )
+    assert (
+        standard_label_uncorrelatedness_scores[f"{test_attribute}_score"]
+        > attribute_filter_label_uncorrelatedness_scores[f"{test_attribute}_score"]
+    )
+
+
+@pytest.mark.parametrize(
+    "test_attribute",
+    [
+        "dark",
+        pytest.param(
+            "blurry",
+            marks=pytest.mark.xfail(
+                reason="odd aspect ratio filter seems to score lower", strict=True
+            ),
+        ),
+        "odd_aspect_ratio",
+    ],
+)
+def test_smallest_scores_with_filters(test_attribute):
+    """
+    Tests that each specific filter has the smallest uncorrelatedness score for its respective property.
+
+    Asserts:
+        AssertionError: If any specific filter score is not the smallest for its respective property.
+    """
+
+    attributes_to_score = ["dark", "blurry", "odd_aspect_ratio"]
+    standard_correlation_scores = get_label_uncorrelatedness_scores()
+
+    score_key = f"{test_attribute}_score"
+    filtered_scores = {
+        f: get_label_uncorrelatedness_scores(circle_filter=f) for f in attributes_to_score
+    }
+
+    # The attribute being tested should have the lowest score for the filtered dataset
+    test_scores = filtered_scores.pop(test_attribute)
+    assert test_scores[score_key] <= min(
+        standard_correlation_scores[score_key],
+        *[scores[score_key] for scores in filtered_scores.values()],
+    )
+
+
+def get_report_text(lab):
+    with contextlib.redirect_stdout(io.StringIO()) as f:
+        lab.report()
+    return f.getvalue()
+
+
+@pytest.mark.parametrize(
+    "test_attribute",
+    [
+        "dark",
+        pytest.param(
+            "blurry",
+            marks=pytest.mark.xfail(
+                reason="blurry filter makes other image properties like 'dark' and 'low information' spurious rather than 'blurry'",
+                strict=False,
+            ),
+        ),
+        "odd_aspect_ratio",
+        "identity",
+    ],
+)
+class TestImagelabReporterAdapter:
+    """
+    Test class for `ImagelabReporterAdapter` to verify the behavior of the `lab.report()` method
+    when handling different image attributes.
+
+    This class uses parameterized testing to check the following:
+
+    1. Output Verification: Ensures that the `report` method prints the expected output based on the `test_attribute` value.
+    2. Spurious Correlations: Confirms that spurious correlations are shown or hidden appropriately:
+    - When `test_attribute` is set to 'identity', the report should not display any spurious correlations.
+    - When `test_attribute` is set to 'dark', 'blurry', or 'odd_aspect_ratio', the report should display the relevant spurious correlations.
+
+    Each test run verifies that the output matches the expected print statements and that the report behaves as expected for each attribute scenario.
+    """
+
+    @pytest.fixture(autouse=True)
+    def lab(self, test_attribute):
+        self.test_attribute = test_attribute
+        self.threshold = 0.01
+        dataset = generate_dataset(circle_filter=test_attribute)
+        lab = Datalab(data=dataset, label_name="label", image_key="image")
+        lab.find_issues()  # Easiest way to get default imagelab checks
+        # Rerun checks with threshold for spurious correlations for simplicity.
+        lab.find_issues(issue_types={"spurious_correlations": {"threshold": self.threshold}})
+        self.label_uncorrelatedness_scores = lab.get_info("spurious_correlations")[
+            "correlations_df"
+        ]
+        return lab
+
+    def _get_correlated_properties(self):
+        if self.label_uncorrelatedness_scores.empty:
+            return []
+        return self.label_uncorrelatedness_scores.query("score < @self.threshold")[
+            "property"
+        ].tolist()
+
+    def _get_label_uncorrelatedness_scores(self):
+        correlated_properties = self._get_correlated_properties()
+        filtered_label_uncorrelatedness_scores = self.label_uncorrelatedness_scores.query(
+            "property in @correlated_properties"
+        )
+        filtered_label_uncorrelatedness_scores.loc[:, "property"] = (
+            filtered_label_uncorrelatedness_scores["property"].apply(
+                lambda x: x.replace("_score", "")
+            )
+        )
+        return filtered_label_uncorrelatedness_scores
+
+    @mock.patch("cleanvision.utils.viz_manager.VizManager.individual_images")
+    def test_report(self, mock_individual_images, lab):
+        report = get_report_text(lab)
+
+        report_correlation_header = "Summary of (potentially spurious) correlations between image properties and class labels detected in the data:\n\n"
+        report_correlation_metric = "Lower scores below correspond to images properties that are more strongly correlated with the class labels.\n\n"
+        filtered_label_uncorrelatedness_scores = self._get_label_uncorrelatedness_scores()
+
+        if self.test_attribute != "identity":
+            assert report_correlation_header in report, "Report should contain correlation header"
+            assert (
+                report_correlation_metric in report
+            ), "Report should contain correlation metric description"
+            assert self.test_attribute in filtered_label_uncorrelatedness_scores["property"].values
+            assert filtered_label_uncorrelatedness_scores.to_string(index=False) in report
+        else:
+            assert (
+                report_correlation_header not in report
+            ), "Report should not contain correlation header"
+            assert (
+                report_correlation_metric not in report
+            ), "Report should not contain correlation metric description"
+            assert filtered_label_uncorrelatedness_scores.empty
+            assert "correlation" not in report.lower()
+            assert "spurious" not in report.lower()
+
+
+@mock.patch("cleanvision.utils.viz_manager.VizManager.individual_images")
+def test_report_image_key(mock_individual_images):
+    X = np.random.rand(100, 2)
+    y = np.sum(X, axis=1)
+    data = {"X": X, "y": y}
+    lab_without_image_key = Datalab(data, label_name="y")
+    lab_without_image_key.find_issues()
+
+    dataset = generate_dataset(circle_filter="dark")
+    lab = Datalab(data=dataset, label_name="label", image_key="image")
+    lab.find_issues()
+
+    report_without_image_key = get_report_text(lab_without_image_key)
+    report = get_report_text(lab)
+
+    report_correlation_header = "Summary of (potentially spurious) correlations between image properties and class labels detected in the data:\n\n"
+    report_correlation_metric = "Lower scores below correspond to images properties that are more strongly correlated with the class labels.\n\n"
+    assert report_correlation_header not in report_without_image_key
+    assert report_correlation_metric not in report_without_image_key
+    assert report_correlation_header in report
+    assert report_correlation_metric in report
+
+    assert "correlation" not in report_without_image_key.lower()
+    assert "spurious" not in report_without_image_key.lower()
+
+
+def test_iterative_property_threshold():
+    dataset = generate_dataset()
+    lab = Datalab(data=dataset, label_name="label", image_key="image")
+
+    lab.find_issues(
+        issue_types={"image_issue_types": {"dark": {}}, "spurious_correlations": {"threshold": 0.1}}
+    )
+    spurious_correlations_info = lab.get_info("spurious_correlations")
+
+    # Spurious correlation scores dataframe should include only 1 row: 'dark_score'
+    assert len(spurious_correlations_info["correlations_df"]) == 1
+    assert spurious_correlations_info["correlations_df"].loc[0, "property"] == "dark_score"
+    # threshold value should be set to 0.1 (user input) for spurious correlations
+    assert spurious_correlations_info["threshold"] == 0.1
+
+    lab.find_issues(
+        issue_types={"image_issue_types": {"dark": {}, "blurry": {}}, "spurious_correlations": {}}
+    )
+    spurious_correlations_info = lab.get_info("spurious_correlations")
+
+    # Spurious correlation scores dataframe should include 2 rows: 'dark_score', 'blurry_score'
+    assert len(spurious_correlations_info["correlations_df"]) == 2
+    assert spurious_correlations_info["correlations_df"]["property"].tolist() == [
+        "dark_score",
+        "blurry_score",
+    ]
+    # threshold value should be set to the default value for spurious correlations
+    default_spurious_correlations_threshold = SPURIOUS_CORRELATION_ISSUE["spurious_correlations"][
+        "threshold"
+    ]
+    assert spurious_correlations_info["threshold"] == default_spurious_correlations_threshold
+
+    lab.find_issues(issue_types={"spurious_correlations": {"threshold": 0.5}})
+    spurious_correlations_info = lab.get_info("spurious_correlations")
+
+    # Spurious correlation scores dataframe should include only 1 row: 'blurry_score'
+    assert len(spurious_correlations_info["correlations_df"]) == 2
+    assert spurious_correlations_info["correlations_df"]["property"].tolist() == [
+        "dark_score",
+        "blurry_score",
+    ]
+    # threshold value should be set to 0.5 (user input) for spurious correlations
+    assert spurious_correlations_info["threshold"] == 0.5
+
+
+def test_get_info_spurious_correlations_check_skipped():
+    """An extra set of tests to verify that spurious correlations
+    are not available without pre-existing image-property checks,
+    so no info on spurious correlations can be retrieved."""
+    dataset = generate_dataset()
+    lab = Datalab(data=dataset, label_name="label", image_key="image")
+
+    # Test 1: Spurious correlations not available without checks
+    error_message = "Spurious correlations have not been calculated. Run find_issues() first."
+    with pytest.raises(ValueError, match=re.escape(error_message)):
+        lab.get_info("spurious_correlations")
+
+    # Test 2: Spurious correlations check skipped without image-property checks
+    with contextlib.redirect_stdout(io.StringIO()) as f:
+        lab.find_issues(issue_types={"spurious_correlations": {}})
+
+    assert f.getvalue().startswith(
+        "Skipping spurious correlations check: Image property scores not available.\n"
+        "To include this check, run find_issues() without parameters to compute all scores."
+    )
+
+    with pytest.raises(ValueError, match=re.escape(error_message)):
+        lab.get_info("spurious_correlations")
+
+    # Test 3: Running image-property checks allows spurious correlations calculation
+    lab.find_issues(issue_types={"image_issue_types": {"odd_aspect_ratio": {}}})
+    lab.find_issues(issue_types={"spurious_correlations": {}})
+
+    spurious_correlations_info = lab.get_info("spurious_correlations")
+    assert set(["correlations_df", "threshold"]) == set(list(spurious_correlations_info.keys()))

tests/test_classification.py

@@ -0,0 +1,991 @@
+from copy import deepcopy
+import sys
+from sklearn.linear_model import LogisticRegression
+from sklearn.base import BaseEstimator
+from sklearn.model_selection import GridSearchCV
+import sklearn
+import scipy
+import pytest
+import numpy as np
+import pandas as pd
+from cleanlab.classification import CleanLearning
+from cleanlab.benchmarking.noise_generation import generate_noise_matrix_from_trace
+from cleanlab.benchmarking.noise_generation import generate_noisy_labels
+from cleanlab.internal.latent_algebra import compute_inv_noise_matrix
+from cleanlab.count import (
+    compute_confident_joint,
+    estimate_cv_predicted_probabilities,
+    get_confident_thresholds,
+)
+from cleanlab.filter import find_label_issues
+
+SEED = 1
+
+
+def make_data(
+    format="numpy",
+    means=[[3, 2], [7, 7], [0, 8]],
+    covs=[[[5, -1.5], [-1.5, 1]], [[1, 0.5], [0.5, 4]], [[5, 1], [1, 5]]],
+    sizes=[100, 50, 50],
+    avg_trace=0.8,
+    seed=SEED,  # set to None for non-reproducible randomness
+):
+    """format specifies what X (and y) looks like, one of:
+    'numpy', 'sparse', 'dataframe', or 'series'.
+    """
+    np.random.seed(seed=seed)
+
+    K = len(means)  # number of classes
+    data = []
+    labels = []
+    test_data = []
+    test_labels = []
+
+    for idx in range(K):
+        data.append(np.random.multivariate_normal(mean=means[idx], cov=covs[idx], size=sizes[idx]))
+        test_data.append(
+            np.random.multivariate_normal(mean=means[idx], cov=covs[idx], size=sizes[idx])
+        )
+        labels.append(np.array([idx for i in range(sizes[idx])]))
+        test_labels.append(np.array([idx for i in range(sizes[idx])]))
+    X_train = np.vstack(data)
+    true_labels_train = np.hstack(labels)
+    X_test = np.vstack(test_data)
+    true_labels_test = np.hstack(test_labels)
+
+    if format == "sparse":
+        X_train = scipy.sparse.csr_matrix(X_train)
+        X_test = scipy.sparse.csr_matrix(X_test)
+    elif format == "dataframe":
+        X_train = pd.DataFrame(X_train)
+        X_test = pd.DataFrame(X_test)
+        # true_labels_train = list(true_labels_train)
+        # true_labels_test = list(true_labels_test)
+    elif format == "series":
+        X_train = pd.Series(X_train[:, 0])
+        X_test = pd.Series(X_test[:, 0])
+        # true_labels_train = pd.Series(true_labels_train)
+        # true_labels_test = pd.Series(true_labels_test)
+    elif format != "numpy":
+        raise ValueError("invalid value specified for: `format`.")
+
+    # Compute p(true_label=k)
+    py = np.bincount(true_labels_train) / float(len(true_labels_train))
+
+    noise_matrix = generate_noise_matrix_from_trace(
+        K,
+        trace=avg_trace * K,
+        py=py,
+        valid_noise_matrix=True,
+        seed=seed,
+    )
+
+    # Generate our noisy labels using the noise_matrix.
+    s = generate_noisy_labels(true_labels_train, noise_matrix)
+    ps = np.bincount(s) / float(len(s))
+
+    return {
+        "X_train": X_train,
+        "true_labels_train": true_labels_train,
+        "X_test": X_test,
+        "true_labels_test": true_labels_test,
+        "labels": s,
+        "ps": ps,
+        "py": py,
+        "noise_matrix": noise_matrix,
+    }
+
+
+def make_rare_label(data):
+    """Makes one label really rare in the dataset."""
+    data = deepcopy(data)
+    y = data["labels"]
+    class0_inds = np.where(y == 0)[0]
+    if len(class0_inds) < 1:
+        raise ValueError("Class 0 too rare already")
+    class0_inds_remove = class0_inds[1:]
+    if len(class0_inds_remove) > 0:
+        y[class0_inds_remove] = 1
+    data["labels"] = y
+    return data
+
+
+def make_high_dim_data(seed=SEED):
+    np.random.seed(seed=seed)
+    X_train = np.random.randint(0, 255, (200, 28, 28))
+    label_train = np.random.randint(0, 10, 200)
+    X_test = np.random.randint(0, 255, (50, 28, 28))
+    label_test = np.random.randint(0, 10, 50)
+    X_train, X_test = X_train / 255.0, X_test / 255.0
+
+    return {
+        "X_train": X_train,
+        "labels_train": label_train,
+        "X_test": X_test,
+        "labels_test": label_test,
+    }
+
+
+DATA = make_data(format="numpy", seed=SEED)
+SPARSE_DATA = make_data(format="sparse", seed=SEED)
+DATAFRAME_DATA = make_data(format="dataframe", seed=SEED)
+SERIES_DATA = make_data(format="series", seed=SEED)  # special case not checked in most tests
+HIGH_DIM_DATA = make_high_dim_data(seed=SEED)
+DATA_FORMATS = {
+    "numpy": DATA,
+    "sparse": SPARSE_DATA,
+    "dataframe": DATAFRAME_DATA,
+}
+
+
+@pytest.mark.parametrize("data", list(DATA_FORMATS.values()))
+def test_cl(data):
+    cl = CleanLearning(clf=LogisticRegression(solver="lbfgs", random_state=SEED))
+    X_train_og = deepcopy(data["X_train"])
+    cl.fit(data["X_train"], data["labels"])
+    score = cl.score(data["X_test"], data["true_labels_test"])
+    print(score)
+    # ensure data has not been altered:
+    if isinstance(X_train_og, np.ndarray):
+        assert (data["X_train"] == X_train_og).all()
+    elif isinstance(X_train_og, pd.DataFrame):
+        assert data["X_train"].equals(X_train_og)
+
+
+def test_cl_default_clf():
+    cl = CleanLearning()  # default clf is LogisticRegression
+    X_train_og = deepcopy(HIGH_DIM_DATA["X_train"])
+    cl.fit(HIGH_DIM_DATA["X_train"], HIGH_DIM_DATA["labels_train"])
+
+    # assert result has the correct length
+    result = cl.predict(HIGH_DIM_DATA["X_test"])
+    assert len(result) == len(HIGH_DIM_DATA["X_test"])
+
+    result = cl.predict(X=HIGH_DIM_DATA["X_test"])
+    assert len(result) == len(HIGH_DIM_DATA["X_test"])
+
+    # assert pred_proba has the right dimensions (N x K),
+    # where K = 10 (number of classes) as specified in make_high_dim_data()
+    pred_proba = cl.predict_proba(HIGH_DIM_DATA["X_test"])
+    assert pred_proba.shape == (len(HIGH_DIM_DATA["X_test"]), 10)
+
+    pred_proba = cl.predict_proba(X=HIGH_DIM_DATA["X_test"])
+    assert pred_proba.shape == (len(HIGH_DIM_DATA["X_test"]), 10)
+
+    score = cl.score(HIGH_DIM_DATA["X_test"], HIGH_DIM_DATA["labels_test"])
+
+    cl.find_label_issues(HIGH_DIM_DATA["X_train"], HIGH_DIM_DATA["labels_train"])
+
+    # ensure data has not been altered:
+    assert (HIGH_DIM_DATA["X_train"] == X_train_og).all()
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+@pytest.mark.parametrize("data", list(DATA_FORMATS.values()))
+def test_rare_label(data):
+    data = make_rare_label(data)
+    test_cl(data)
+
+
+def test_invalid_inputs():
+    data = make_data(sizes=[1, 1, 1])
+    try:
+        test_cl(data)
+    except Exception as e:
+        assert "Need more data" in str(e)
+    else:
+        raise Exception("expected test to raise Exception")
+    try:
+        cl = CleanLearning(
+            clf=LogisticRegression(solver="lbfgs", random_state=SEED),
+            find_label_issues_kwargs={"return_indices_ranked_by": "self_confidence"},
+        )
+        cl.fit(
+            data["X_train"],
+            data["labels"],
+        )
+    except Exception as e:
+        assert "not supported" in str(e) or "Need more data from each class" in str(e)
+    else:
+        raise Exception("expected test to raise Exception")
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_aux_inputs():
+    data = DATA
+    K = len(np.unique(data["labels"]))
+    confident_joint = np.ones(shape=(K, K))
+    np.fill_diagonal(confident_joint, 10)
+    find_label_issues_kwargs = {
+        "confident_joint": confident_joint,
+        "min_examples_per_class": 2,
+    }
+    cl = CleanLearning(
+        clf=LogisticRegression(solver="lbfgs", random_state=SEED),
+        find_label_issues_kwargs=find_label_issues_kwargs,
+        verbose=1,
+    )
+    label_issues_df = cl.find_label_issues(data["X_train"], data["labels"], clf_kwargs={})
+    assert isinstance(label_issues_df, pd.DataFrame)
+    FIND_OUTPUT_COLUMNS = ["is_label_issue", "label_quality", "given_label", "predicted_label"]
+    assert list(label_issues_df.columns) == FIND_OUTPUT_COLUMNS
+    assert label_issues_df.equals(cl.get_label_issues())
+    cl.fit(
+        data["X_train"],
+        data["labels"],
+        label_issues=label_issues_df,
+        clf_kwargs={},
+        clf_final_kwargs={},
+    )
+    label_issues_df = cl.get_label_issues()
+    assert isinstance(label_issues_df, pd.DataFrame)
+    assert list(label_issues_df.columns) == (FIND_OUTPUT_COLUMNS + ["sample_weight"])
+    score = cl.score(data["X_test"], data["true_labels_test"])
+
+    # Test a second fit
+    cl.fit(data["X_train"], data["labels"])
+
+    # Test cl.find_label_issues with pred_prob input
+    pred_probs_test = cl.predict_proba(data["X_test"])
+    label_issues_df = cl.find_label_issues(
+        X=None, labels=data["true_labels_test"], pred_probs=pred_probs_test
+    )
+    assert isinstance(label_issues_df, pd.DataFrame)
+    assert list(label_issues_df.columns) == FIND_OUTPUT_COLUMNS
+    assert label_issues_df.equals(cl.get_label_issues())
+    cl.save_space()
+    assert cl.label_issues_df is None
+
+    # Verbose off
+    cl = CleanLearning(clf=LogisticRegression(solver="lbfgs", random_state=SEED), verbose=0)
+    cl.save_space()  # dummy call test
+
+    cl = CleanLearning(clf=LogisticRegression(solver="lbfgs", random_state=SEED), verbose=0)
+    cl.find_label_issues(
+        labels=data["true_labels_test"], pred_probs=pred_probs_test, save_space=True
+    )
+
+    cl = CleanLearning(clf=LogisticRegression(solver="lbfgs", random_state=SEED), verbose=1)
+
+    # Test with label_issues_mask input
+    label_issues_mask = find_label_issues(
+        labels=data["true_labels_test"],
+        pred_probs=pred_probs_test,
+    )
+    cl.fit(data["X_test"], data["true_labels_test"], label_issues=label_issues_mask)
+    label_issues_df = cl.get_label_issues()
+    assert isinstance(label_issues_df, pd.DataFrame)
+    assert set(label_issues_df.columns).issubset(FIND_OUTPUT_COLUMNS)
+
+    # Test with label_issues_indices input
+    label_issues_indices = find_label_issues(
+        labels=data["true_labels_test"],
+        pred_probs=pred_probs_test,
+        return_indices_ranked_by="confidence_weighted_entropy",
+    )
+    cl.fit(data["X_test"], data["true_labels_test"], label_issues=label_issues_indices)
+    label_issues_df2 = cl.get_label_issues().copy()
+    assert isinstance(label_issues_df2, pd.DataFrame)
+    assert set(label_issues_df2.columns).issubset(FIND_OUTPUT_COLUMNS)
+    assert label_issues_df2["is_label_issue"].equals(label_issues_df["is_label_issue"])
+
+    # Test fit() with pred_prob input:
+    cl.fit(
+        data["X_test"],
+        data["true_labels_test"],
+        pred_probs=pred_probs_test,
+        label_issues=label_issues_mask,
+    )
+    label_issues_df = cl.get_label_issues()
+    assert isinstance(label_issues_df, pd.DataFrame)
+    assert set(label_issues_df.columns).issubset(FIND_OUTPUT_COLUMNS)
+    assert "label_quality" in label_issues_df.columns
+
+    # Test with sample_weight input:
+    cl = CleanLearning(clf=LogisticRegression(solver="lbfgs", random_state=SEED), verbose=1)
+    cl.fit(
+        data["X_test"],
+        data["true_labels_test"],
+        sample_weight=np.random.randn(len(data["true_labels_test"])),
+    )
+    cl.fit(
+        data["X_test"],
+        data["true_labels_test"],
+        label_issues=cl.get_label_issues(),
+        sample_weight=np.random.randn(len(data["true_labels_test"])),
+    )
+
+
+class LogisticRegressionWithValidationData(LogisticRegression):
+    def fit(self, X, y, X_val=None, y_val=None):
+        super().fit(X, y)
+
+        # Final fit() call does not use validation data
+        # Checks to prevent arg missing error
+        if X_val is not None or y_val is not None:
+            print(self.score(X_val, y_val))
+
+
+def val_func(X_val, y_val):
+    return {"X_val": X_val, "y_val": y_val}
+
+
+def test_validation_data():
+    data = DATA
+    cl = CleanLearning(clf=LogisticRegressionWithValidationData())
+    cl.fit(
+        data["X_train"],
+        data["labels"],
+        validation_func=val_func,
+    )
+
+
+def test_raise_error_no_clf_fit():
+    class struct(object):
+        def predict(self):
+            pass
+
+        def predict_proba(self):
+            pass
+
+    try:
+        CleanLearning(clf=struct())
+    except Exception as e:
+        assert "fit" in str(e)
+        with pytest.raises(ValueError) as e:
+            CleanLearning(clf=struct())
+
+
+def test_raise_error_no_clf_predict_proba():
+    class struct(object):
+        def fit(self):
+            pass
+
+        def predict(self):
+            pass
+
+    try:
+        CleanLearning(clf=struct())
+    except Exception as e:
+        assert "predict_proba" in str(e)
+        with pytest.raises(ValueError) as e:
+            CleanLearning(clf=struct())
+
+
+def test_raise_error_no_clf_predict():
+    class struct(object):
+        def fit(self):
+            pass
+
+        def predict_proba(self):
+            pass
+
+    try:
+        CleanLearning(clf=struct())
+    except Exception as e:
+        assert "predict" in str(e)
+        with pytest.raises(ValueError) as e:
+            CleanLearning(clf=struct())
+
+
+def test_seed():
+    cl = CleanLearning(seed=SEED)
+    assert cl.seed is not None
+
+
+def test_default_clf():
+    cl = CleanLearning()
+    check1 = cl.clf is not None and hasattr(cl.clf, "fit")
+    check2 = hasattr(cl.clf, "predict") and hasattr(cl.clf, "predict_proba")
+    assert check1 and check2
+
+
+def test_clf_fit_nm():
+    cl = CleanLearning()
+    # Example of a bad noise matrix (impossible to learn from)
+    nm = np.array([[0, 1], [1, 0]])
+    try:
+        cl.fit(X=np.arange(3), labels=np.array([0, 0, 1]), noise_matrix=nm)
+    except Exception as e:
+        assert "Trace(noise_matrix)" in str(e)
+        with pytest.raises(ValueError) as e:
+            cl.fit(X=np.arange(3), labels=np.array([0, 0, 1]), noise_matrix=nm)
+
+
+def test_clf_fit_inm():
+    cl = CleanLearning()
+    # Example of a bad noise matrix (impossible to learn from)
+    inm = np.array([[0.1, 0.9], [0.9, 0.1]])
+    try:
+        cl.fit(X=np.arange(3), labels=np.array([0, 0, 1]), inverse_noise_matrix=inm)
+    except Exception as e:
+        assert "Trace(inverse_noise_matrix)" in str(e)
+        with pytest.raises(ValueError) as e:
+            cl.fit(X=np.arange(3), labels=np.array([0, 0, 1]), inverse_noise_matrix=inm)
+
+
+@pytest.mark.parametrize("format", list(DATA_FORMATS.keys()))
+def test_fit_with_nm(
+    format,
+    seed=SEED,
+    used_by_another_test=False,
+):
+    data = DATA_FORMATS[format]
+    cl = CleanLearning(
+        seed=seed,
+    )
+    nm = data["noise_matrix"]
+    # Learn with noisy labels with noise matrix given
+    cl.fit(data["X_train"], data["labels"], noise_matrix=nm)
+    score_nm = cl.score(data["X_test"], data["true_labels_test"])
+    # Learn with noisy labels and estimate the noise matrix.
+    cl2 = CleanLearning(
+        seed=seed,
+    )
+    cl2.fit(
+        data["X_train"],
+        data["labels"],
+    )
+    score = cl2.score(data["X_test"], data["true_labels_test"])
+    if used_by_another_test:
+        return score, score_nm
+    else:
+        assert score < score_nm + 1e-4
+
+
+@pytest.mark.parametrize("format", list(DATA_FORMATS.keys()))
+def test_fit_with_inm(
+    format,
+    seed=SEED,
+    used_by_another_test=False,
+):
+    data = DATA_FORMATS[format]
+    cl = CleanLearning(
+        seed=seed,
+    )
+    inm = compute_inv_noise_matrix(
+        py=data["py"],
+        noise_matrix=data["noise_matrix"],
+        ps=data["ps"],
+    )
+    # Learn with noisy labels with inverse noise matrix given
+    cl.fit(data["X_train"], data["labels"], inverse_noise_matrix=inm)
+    score_inm = cl.score(data["X_test"], data["true_labels_test"])
+    # Learn with noisy labels and estimate the inv noise matrix.
+    cl2 = CleanLearning(
+        seed=seed,
+    )
+    cl2.fit(
+        data["X_train"],
+        data["labels"],
+    )
+    score = cl2.score(data["X_test"], data["true_labels_test"])
+    if used_by_another_test:
+        return score, score_inm
+    else:
+        assert score < score_inm + 1e-4
+
+
+@pytest.mark.parametrize("format", list(DATA_FORMATS.keys()))
+def test_clf_fit_nm_inm(format):
+    data = DATA_FORMATS[format]
+    cl = CleanLearning(seed=SEED)
+    nm = data["noise_matrix"]
+    inm = compute_inv_noise_matrix(
+        py=data["py"],
+        noise_matrix=nm,
+        ps=data["ps"],
+    )
+    cl.fit(
+        X=data["X_train"],
+        labels=data["labels"],
+        noise_matrix=nm,
+        inverse_noise_matrix=inm,
+    )
+    score_nm_inm = cl.score(data["X_test"], data["true_labels_test"])
+
+    # Learn with noisy labels and estimate the inv noise matrix.
+    cl2 = CleanLearning(seed=SEED)
+    cl2.fit(
+        data["X_train"],
+        data["labels"],
+    )
+    score = cl2.score(data["X_test"], data["true_labels_test"])
+    assert score < score_nm_inm + 1e-4
+
+
+@pytest.mark.parametrize("format", list(DATA_FORMATS.keys()))
+def test_clf_fit_y_alias(format):
+    data = DATA_FORMATS[format]
+    cl = CleanLearning(seed=SEED)
+
+    # Valid signature
+    cl.fit(data["X_train"], data["labels"])
+
+    # Valid signature for labels/y alias
+    cl.fit(data["X_train"], labels=data["labels"])
+    cl.fit(data["X_train"], y=data["labels"])
+    cl.fit(X=data["X_train"], labels=data["labels"])
+    cl.fit(X=data["X_train"], y=data["labels"])
+
+    # Invalid signatures
+    with pytest.raises(ValueError):
+        cl.fit(data["X_train"])
+    with pytest.raises(ValueError):
+        cl.fit(data["X_train"], data["labels"], y=data["labels"])
+    with pytest.raises(ValueError):
+        cl.fit(X=data["X_train"], labels=data["labels"], y=data["labels"])
+
+
+@pytest.mark.parametrize("format", list(DATA_FORMATS.keys()))
+def test_pred_and_pred_proba(format):
+    data = DATA_FORMATS[format]
+    cl = CleanLearning()
+    cl.fit(data["X_train"], data["labels"])
+    n = np.shape(data["true_labels_test"])[0]
+    m = len(np.unique(data["true_labels_test"]))
+    pred = cl.predict(data["X_test"])
+    probs = cl.predict_proba(data["X_test"])
+    # Just check that this functions return what we expect
+    assert np.shape(pred)[0] == n
+    assert np.shape(probs) == (n, m)
+
+
+@pytest.mark.parametrize("format", list(DATA_FORMATS.keys()))
+def test_score(format):
+    data = DATA_FORMATS[format]
+    phrase = "cleanlab is dope"
+
+    class Struct:
+        def fit(self):
+            pass
+
+        def predict_proba(self):
+            pass
+
+        def predict(self):
+            pass
+
+        def score(self, X, y):
+            return phrase
+
+    cl = CleanLearning(clf=Struct())
+    score = cl.score(data["X_test"], data["true_labels_test"])
+    assert score == phrase
+
+
+@pytest.mark.parametrize("format", list(DATA_FORMATS.keys()))
+def test_no_score(format):
+    data = DATA_FORMATS[format]
+
+    class Struct:
+        def fit(self):
+            pass
+
+        def predict_proba(self):
+            pass
+
+        def predict(self, X):
+            return data["true_labels_test"]
+
+    cl = CleanLearning(clf=Struct())
+    score = cl.score(data["X_test"], data["true_labels_test"])
+    assert abs(score - 1) < 1e-6
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+@pytest.mark.parametrize("format", list(DATA_FORMATS.keys()))
+def test_no_fit_sample_weight(format):
+    data = DATA_FORMATS[format]
+
+    class Struct:
+        def fit(self, X, y):
+            pass
+
+        def predict_proba(self, X):
+            n_samples = len(X)
+            n_classes = len(np.unique(data["true_labels_train"]))
+            return np.ones((n_samples, n_classes)) / n_classes
+
+        def predict(self, X):
+            return np.zeros(len(X), dtype=int)
+
+    n = np.shape(data["true_labels_test"])[0]
+    m = len(np.unique(data["true_labels_test"]))
+    pred_probs = np.ones((n, m)) / m
+    cl = CleanLearning(clf=Struct())
+    cl.fit(
+        data["X_train"],
+        data["true_labels_train"],
+        pred_probs=pred_probs,
+        noise_matrix=data["noise_matrix"],
+    )
+    # If we make it here, without any error:
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+@pytest.mark.parametrize("format", list(DATA_FORMATS.keys()))
+def test_fit_pred_probs(format):
+    data = DATA_FORMATS[format]
+
+    cl = CleanLearning()
+    pred_probs = estimate_cv_predicted_probabilities(
+        X=data["X_train"],
+        labels=data["true_labels_train"],
+    )
+    cl.fit(X=data["X_train"], labels=data["true_labels_train"], pred_probs=pred_probs)
+    score_with_pred_probs = cl.score(data["X_test"], data["true_labels_test"])
+    cl = CleanLearning()
+    cl.fit(
+        X=data["X_train"],
+        labels=data["true_labels_train"],
+    )
+    score_no_pred_probs = cl.score(data["X_test"], data["true_labels_test"])
+    assert abs(score_with_pred_probs - score_no_pred_probs) < 0.01
+
+
+def make_2d(X):
+    X = np.asarray(X)
+    return X.reshape(X.shape[0], -1)
+
+
+class ReshapingLogisticRegression(BaseEstimator):
+    def __init__(self):
+        self.clf = LogisticRegression()
+
+    def fit(self, X, y):
+        y = np.asarray(y).flatten()
+        self.clf.fit(make_2d(X), y)
+
+    def predict(self, X):
+        return self.clf.predict(make_2d(X))
+
+    def predict_proba(self, X):
+        return self.clf.predict_proba(make_2d(X))
+
+    def score(self, X, y, sample_weight=None):
+        return self.clf.score(make_2d(X), y, sample_weight=sample_weight)
+
+
+def dimN_data(N):
+    size = [100] + [3 for _ in range(N - 1)]
+    X = np.random.normal(size=size)
+    labels = np.random.randint(0, 4, size=100)
+    # ensure that every class is represented
+    labels[0:10] = 0
+    labels[11:20] = 1
+    labels[21:30] = 2
+    labels[31:40] = 3
+    return X, labels
+
+
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+@pytest.mark.parametrize("N", [1, 3, 4])
+def test_dimN(N):
+    X, labels = dimN_data(N)
+    cl = CleanLearning(clf=ReshapingLogisticRegression())
+    # just make sure we don't crash...
+    cl.fit(X, labels)
+    cl.predict(X)
+    cl.predict_proba(X)
+    cl.score(X, labels)
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_1D_formats():
+    X, labels = dimN_data(1)
+    X_series = pd.Series(X)
+    labels_series = pd.Series(labels)
+    idx = list(np.random.choice(len(labels), size=len(labels), replace=False))
+    X_series.index = idx
+    labels_series.index = idx
+    cl = CleanLearning(clf=ReshapingLogisticRegression())
+    # just make sure we don't crash...
+    cl.fit(X_series, labels_series)
+    cl.predict(X_series)
+    cl.predict_proba(X_series)
+    cl.score(X_series, labels)
+    # Repeat with rare labels:
+    labels_rare = deepcopy(labels)
+    class0_inds = np.where(labels_rare == 0)[0]
+    class0_inds_remove = class0_inds[1:]
+    labels_rare[class0_inds_remove] = 1
+    cl = CleanLearning(clf=ReshapingLogisticRegression())
+    cl.fit(X_series, labels_rare)
+    cl.predict(X_series)
+    cl.predict_proba(X_series)
+    cl.score(X_series, labels)
+    # Repeat with DataFrame labels:
+    labels_df = pd.DataFrame({"colname": labels})
+    cl = CleanLearning(clf=ReshapingLogisticRegression())
+    cl.fit(X, labels_df)
+    cl.predict(X)
+    pred_probs = cl.predict_proba(X)
+    cl.score(X, labels)
+    # Repeat with DataFrame labels and pred_probs
+    cl = CleanLearning(clf=ReshapingLogisticRegression())
+    cl.fit(X, labels_df, pred_probs=pred_probs)
+    # Repeat with list labels:
+    labels_list = list(labels)
+    cl = CleanLearning(clf=ReshapingLogisticRegression())
+    cl.fit(X, labels_list)
+    cl.predict(X)
+    cl.predict_proba(X)
+    cl.score(X, labels)
+
+
+# Check if the current Python version is 3.11
+is_python_311 = sys.version_info.major == 3 and sys.version_info.minor == 11
+
+# This warning should be ignored as in Python 3.11, the sre_constants module has been deprecated.
+# At the time of writing this, cleanlab supports Python 3.8-3.11. This warning is raised by
+# tensorflow <2.14.0, which imports sre_constants. This warning is not relevant to cleanlab.
+# Once Python 3.8 reaches EOL, we may remove this warning filter as we can set the tensorflow
+# dev-dependency to a version that does not raise this warning (2.14 or higher).
+if is_python_311:
+    sre_deprecation_pytestmark = pytest.mark.filterwarnings(
+        "ignore:module 'sre_constants' is deprecated"
+    )
+else:
+    sre_deprecation_pytestmark = pytest.mark.filterwarnings("default")
+
+# Check if the installed version of sklearn is 1.5.0.
+# The test_sklearn_gridsearchcv test fails due to a regression introduced in 1.5.0.
+# This issue will be fixed in sklearn version 1.5.1.
+uses_sklearn_1_5_0 = sklearn.__version__ == "1.5.0"
+
+
+@sre_deprecation_pytestmark  # Allow sre_constants deprecation warning for Python 3.11
+@pytest.mark.filterwarnings("error")  # All other warnings are treated as errors
+@pytest.mark.skipif(
+    uses_sklearn_1_5_0,
+    reason="Test is skipped because sklearn 1.5.0 is installed, which has a regression for GridSearchCV.",
+)  # TODO: Remove this line once sklearn 1.5.1 is released
+def test_sklearn_gridsearchcv():
+    # hyper-parameters for grid search
+    param_grid = {
+        "find_label_issues_kwargs": [
+            {"filter_by": "prune_by_noise_rate"},
+            {"filter_by": "prune_by_class"},
+            {"filter_by": "both"},
+            {"filter_by": "confident_learning"},
+            {"filter_by": "predicted_neq_given"},
+        ],
+        "converge_latent_estimates": [True, False],
+    }
+
+    clf = LogisticRegression(random_state=0, solver="lbfgs")
+
+    cv = GridSearchCV(
+        estimator=CleanLearning(clf),
+        param_grid=param_grid,
+        cv=3,
+    )
+
+    # cv.fit() raises a warning if some fits fail (including raising
+    # exceptions); we don't expect any fits to fail, so ensure that the code
+    # doesn't raise any warnings
+    cv.fit(X=DATA["X_train"], y=DATA["labels"])
+
+
+@pytest.mark.parametrize("filter_by", ["both", "confident_learning"])
+@pytest.mark.parametrize("seed", [0, 6, 2])
+def test_cj_in_find_label_issues_kwargs(filter_by, seed):
+    labels = DATA["labels"]
+    num_issues = []
+    for provide_confident_joint in [True, False]:
+        print(f"\nfilter_by: {filter_by} | seed: {seed} | cj_provided: {provide_confident_joint}")
+        np.random.seed(seed=seed)
+        if provide_confident_joint:
+            pred_probs = estimate_cv_predicted_probabilities(
+                X=DATA["X_train"], labels=labels, seed=seed
+            )
+            confident_joint = compute_confident_joint(labels=labels, pred_probs=pred_probs)
+            cl = CleanLearning(
+                find_label_issues_kwargs={
+                    "confident_joint": confident_joint,
+                    "filter_by": "both",
+                    "min_examples_per_class": 1,
+                },
+                verbose=1,
+            )
+        else:
+            cl = CleanLearning(
+                clf=LogisticRegression(random_state=seed),
+                find_label_issues_kwargs={
+                    "filter_by": "both",
+                    "min_examples_per_class": 1,
+                },
+                verbose=0,
+            )
+        label_issues_df = cl.find_label_issues(DATA["X_train"], labels=labels)
+        label_issues_mask = label_issues_df["is_label_issue"].values
+        # Check if the noise matrix was computed based on the passed in confident joint
+        cj_reconstruct = (cl.inverse_noise_matrix * np.bincount(DATA["labels"])).T.astype(int)
+        np.all(cl.confident_joint == cj_reconstruct)
+        num_issues.append(sum(label_issues_mask))
+
+    # Chceck that the same exact number of issues are found regardless if the confident joint
+    # is computed during find_label_issues or precomputed and provided as a kwargs parameter.
+    assert num_issues[0] == num_issues[1]
+
+
+def test_find_label_issues_uses_thresholds():
+    X = DATA["X_train"]
+    labels = DATA["labels"]
+    pred_probs = estimate_cv_predicted_probabilities(X=X, labels=labels)
+
+    confident_thresholds = get_confident_thresholds(labels=labels, pred_probs=pred_probs)
+    confident_joint = compute_confident_joint(labels=labels, pred_probs=pred_probs)
+
+    # regular find label issues with no args
+    cl = CleanLearning()
+    label_issues_reg = cl.find_label_issues(labels=labels, pred_probs=pred_probs)
+
+    # find label issues with specified confident thresholds
+    cl = CleanLearning()
+    label_issues_thres = cl.find_label_issues(
+        labels=labels, pred_probs=pred_probs, thresholds=confident_thresholds
+    )
+
+    # find label issues with specified confident joint
+    cl = CleanLearning(
+        find_label_issues_kwargs={
+            "confident_joint": confident_joint,
+        }
+    )
+    label_issues_cj = cl.find_label_issues(labels=labels, pred_probs=pred_probs)
+
+    # the labels issues in above three calls should be the same
+    assert np.sum(label_issues_reg["is_label_issue"]) == np.sum(
+        label_issues_thres["is_label_issue"]
+    )
+    assert np.sum(label_issues_reg["is_label_issue"]) == np.sum(label_issues_cj["is_label_issue"])
+
+    # find label issues with different specified confident thresholds
+    confident_thresholds_alt = np.full(pred_probs.shape[1], 0.25)
+    cl = CleanLearning()
+    label_issues_thres_alt = cl.find_label_issues(
+        labels=labels, pred_probs=pred_probs, thresholds=confident_thresholds_alt
+    )
+
+    # find label issues with different specified confident joint
+    confident_joint_alt = compute_confident_joint(
+        labels=labels, pred_probs=pred_probs, thresholds=confident_thresholds_alt
+    )
+    cl = CleanLearning(
+        find_label_issues_kwargs={
+            "confident_joint": confident_joint_alt,
+        }
+    )
+    label_issues_cj_alt = cl.find_label_issues(labels=labels, pred_probs=pred_probs)
+
+    # the number of issues for these 2 alt calls should be same as one another, but different from above 3
+    assert np.sum(label_issues_thres_alt["is_label_issue"]) == np.sum(
+        label_issues_cj_alt["is_label_issue"]
+    )
+    assert np.sum(label_issues_thres_alt["is_label_issue"]) != np.sum(
+        label_issues_reg["is_label_issue"]
+    )
+
+
+def test_find_issues_missing_classes():
+    labels = np.array([0, 0, 2, 2])
+    pred_probs = np.array(
+        [[0.9, 0.0, 0.1, 0.0], [0.8, 0.0, 0.2, 0.0], [0.1, 0.0, 0.9, 0.0], [0.95, 0.0, 0.05, 0.0]]
+    )
+    issues_df = CleanLearning(
+        find_label_issues_kwargs={"min_examples_per_class": 0}
+    ).find_label_issues(labels=labels, pred_probs=pred_probs)
+    issues = issues_df["is_label_issue"].values
+    assert np.all(issues == np.array([False, False, False, True]))
+    # Check results match without these missing classes present in pred_probs:
+    pred_probs2 = pred_probs[:, list(sorted(np.unique(labels)))]
+    labels2 = np.array([0, 0, 1, 1])
+    issues_df2 = CleanLearning(
+        find_label_issues_kwargs={"min_examples_per_class": 0}
+    ).find_label_issues(labels=labels2, pred_probs=pred_probs2)
+    assert all(issues_df2["is_label_issue"].values == issues)
+
+
+def test_find_issues_low_memory():
+    X = DATA["X_train"]
+    labels = DATA["labels"]
+    pred_probs = estimate_cv_predicted_probabilities(X=X, labels=labels, seed=SEED)
+    issues_df = CleanLearning().find_label_issues(labels=labels, pred_probs=pred_probs)
+    issues_df_lm = CleanLearning(low_memory=True).find_label_issues(
+        labels=labels, pred_probs=pred_probs
+    )
+    # check jaccard similarity:
+    intersection = len(list(set(issues_df).intersection(set(issues_df_lm))))
+    union = len(set(issues_df)) + len(set(issues_df_lm)) - intersection
+    assert float(intersection) / union > 0.95
+    # Without pred_probs
+    issues_df = CleanLearning(low_memory=True, verbose=True, seed=SEED).find_label_issues(
+        X=X, labels=labels
+    )
+    assert issues_df.equals(issues_df_lm)
+    # With unused arguments find_label_issues_kwargs and noise_matrix
+    find_label_issues_kwargs = {"min_examples_per_class": 2}
+    issues_df = CleanLearning(
+        low_memory=True, find_label_issues_kwargs=find_label_issues_kwargs, seed=SEED
+    ).find_label_issues(X=X, labels=labels, noise_matrix=DATA["noise_matrix"])
+    assert issues_df.equals(issues_df_lm)
+
+
+def test_confident_joint_setting_in_find_label_issues_kwargs():
+    """
+    This test ensures that the 'confident_joint' is correctly set in the
+    'find_label_issues_kwargs' of the 'CleanLearning' class when calling find_label_issues().
+
+    This test was added to cover the lines of code that were previously
+    missed due to the removal of another test.
+    """
+    # Load training data and labels
+    X = DATA["X_train"]
+    labels = DATA["labels"]
+
+    # Estimate predicted probabilities using cross-validation
+    pred_probs = estimate_cv_predicted_probabilities(X=X, labels=labels, seed=SEED)
+
+    # Initialize CleanLearning instance
+    cl = CleanLearning()
+
+    # Test that the confident joint is not set initially
+    cj = cl.find_label_issues_kwargs.get("confident_joint")
+    assert cj is None, "Initial confident_joint should be None"
+
+    # Call find_label_issues to set the confident joint
+    cl.find_label_issues(labels=labels, pred_probs=pred_probs)
+    cj = cl.find_label_issues_kwargs.get("confident_joint")
+
+    # Compute expected confident joint
+    expected_cj = compute_confident_joint(labels=labels, pred_probs=pred_probs)
+
+    # Assert that the confident joint is set correctly
+    np.testing.assert_array_equal(
+        cj, expected_cj, "Confident joint not set correctly after find_label_issues"
+    )
+
+    # Pass a precomputed confident_joint to the CleanLearning instance
+    cj_as_input = np.random.rand(3, 3)
+    cl = CleanLearning(
+        find_label_issues_kwargs={
+            "confident_joint": cj_as_input,
+        }
+    )
+
+    # Ensure the precomputed confident joint is used
+    cj = cl.find_label_issues_kwargs.get("confident_joint")
+    np.testing.assert_array_equal(
+        cj, cj_as_input, "Confident joint not set correctly when passed as input"
+    )
+
+    # Calling find_label_issues should not change the precomputed confident_joint
+    cl.find_label_issues(labels=labels, pred_probs=pred_probs)
+    cj = cl.find_label_issues_kwargs.get("confident_joint")
+    np.testing.assert_array_equal(
+        cj,
+        cj_as_input,
+        "Confident joint should not change after find_label_issues call when precomputed joint is provided",
+    )

tests/test_data_valuation.py

@@ -0,0 +1,95 @@
+import numpy as np
+import pytest
+from hypothesis import given, settings, strategies as st
+from hypothesis.strategies import composite
+from hypothesis.extra.numpy import arrays
+
+from sklearn.neighbors import NearestNeighbors
+
+from cleanlab.data_valuation import _knn_shapley_score, data_shapley_knn
+from cleanlab.internal.neighbor.knn_graph import create_knn_graph_and_index
+
+
+class TestDataValuation:
+    K = 3
+    N = 100
+    num_features = 10
+
+    @pytest.fixture
+    def features(self):
+        return np.random.rand(self.N, self.num_features)
+
+    @pytest.fixture
+    def labels(self):
+        return np.random.randint(0, 2, self.N)
+
+    @pytest.fixture
+    def knn_graph(self, features):
+        knn = NearestNeighbors(n_neighbors=self.K).fit(features)
+        knn_graph = knn.kneighbors_graph(mode="distance")
+        return knn_graph
+
+    def test_data_shapley_knn(self, labels, features):
+        shapley = data_shapley_knn(labels, features=features, k=self.K)
+        assert shapley.shape == (100,)
+        assert np.all(shapley >= 0)
+        assert np.all(shapley <= 1)
+
+    def test_data_shapley_knn_with_knn_graph(self, labels, knn_graph):
+        shapley = data_shapley_knn(labels, knn_graph=knn_graph, k=self.K)
+        assert shapley.shape == (100,)
+        assert np.all(shapley >= 0)
+        assert np.all(shapley <= 1)
+
+
+@composite
+def valid_data(draw):
+    """
+    A custom strategy to generate valid labels, features, and k such that:
+    - labels and features have the same length
+    - k is less than the length of labels and features
+    """
+    # Generate a valid length for labels and features
+    length = draw(st.integers(min_value=11, max_value=1000))
+
+    # Generate labels and features of the same length
+    labels = draw(
+        arrays(
+            dtype=np.int32,
+            shape=length,
+            elements=st.integers(min_value=0, max_value=length - 1),
+        )
+    )
+    features = draw(
+        arrays(
+            dtype=np.float64,
+            shape=(length, draw(st.integers(min_value=2, max_value=50))),
+            elements=st.floats(min_value=-1.0, max_value=1.0),
+        )
+    )
+
+    # Generate k such that it is less than the length of labels and features
+    k = draw(st.integers(min_value=1, max_value=length - 1))
+
+    return labels, features, k
+
+
+class TestDataShapleyKNNScore:
+    """This test class prioritizes testing the raw/untransformed outputs of the _knn_shapley_score function."""
+
+    @settings(
+        max_examples=1000, deadline=None
+    )  # Increase the number of examples to test more cases
+    @given(valid_data())
+    def test_knn_shapley_score_property(self, data):
+        labels, features, k = data
+
+        knn_graph, _ = create_knn_graph_and_index(features, n_neighbors=k)
+        neighbor_indices = knn_graph.indices.reshape(-1, k)
+
+        scores = _knn_shapley_score(neighbor_indices, labels, k)
+
+        # Shapley scores should be between -1 and 1
+        assert scores.shape == (len(labels),)
+        assert np.all(scores >= -1)
+        assert np.all(scores <= 1)

tests/test_dataset.py

@@ -0,0 +1,564 @@
+import requests
+import pytest
+import hypothesis.extra.numpy as npst
+import hypothesis.strategies as st
+import io
+import numpy as np
+from hypothesis import given, settings
+from cleanlab.dataset import (
+    health_summary,
+    find_overlapping_classes,
+    rank_classes_by_label_quality,
+    overall_label_health_score,
+)
+from cleanlab.count import estimate_joint, num_label_issues, compute_confident_joint
+
+cifar100 = [
+    "apple",
+    "aquarium_fish",
+    "baby",
+    "bear",
+    "beaver",
+    "bed",
+    "bee",
+    "beetle",
+    "bicycle",
+    "bottle",
+    "bowl",
+    "boy",
+    "bridge",
+    "bus",
+    "butterfly",
+    "camel",
+    "can",
+    "castle",
+    "caterpillar",
+    "cattle",
+    "chair",
+    "chimpanzee",
+    "clock",
+    "cloud",
+    "cockroach",
+    "couch",
+    "crab",
+    "crocodile",
+    "cup",
+    "dinosaur",
+    "dolphin",
+    "elephant",
+    "flatfish",
+    "forest",
+    "fox",
+    "girl",
+    "hamster",
+    "house",
+    "kangaroo",
+    "keyboard",
+    "lamp",
+    "lawn_mower",
+    "leopard",
+    "lion",
+    "lizard",
+    "lobster",
+    "man",
+    "maple_tree",
+    "motorcycle",
+    "mountain",
+    "mouse",
+    "mushroom",
+    "oak_tree",
+    "orange",
+    "orchid",
+    "otter",
+    "palm_tree",
+    "pear",
+    "pickup_truck",
+    "pine_tree",
+    "plain",
+    "plate",
+    "poppy",
+    "porcupine",
+    "possum",
+    "rabbit",
+    "raccoon",
+    "ray",
+    "road",
+    "rocket",
+    "rose",
+    "sea",
+    "seal",
+    "shark",
+    "shrew",
+    "skunk",
+    "skyscraper",
+    "snail",
+    "snake",
+    "spider",
+    "squirrel",
+    "streetcar",
+    "sunflower",
+    "sweet_pepper",
+    "table",
+    "tank",
+    "telephone",
+    "television",
+    "tiger",
+    "tractor",
+    "train",
+    "trout",
+    "tulip",
+    "turtle",
+    "wardrobe",
+    "whale",
+    "willow_tree",
+    "wolf",
+    "woman",
+    "worm",
+]
+caltech256 = [
+    "ak47",
+    "american-flag",
+    "backpack",
+    "baseball-bat",
+    "baseball-glove",
+    "basketball-hoop",
+    "bat",
+    "bathtub",
+    "bear",
+    "beer-mug",
+    "billiards",
+    "binoculars",
+    "birdbath",
+    "blimp",
+    "bonsai",
+    "boom-box",
+    "bowling-ball",
+    "bowling-pin",
+    "boxing-glove",
+    "brain",
+    "breadmaker",
+    "buddha",
+    "bulldozer",
+    "butterfly",
+    "cactus",
+    "cake",
+    "calculator",
+    "camel",
+    "cannon",
+    "canoe",
+    "car-tire",
+    "cartman",
+    "cd",
+    "centipede",
+    "cereal-box",
+    "chandelier",
+    "chess-board",
+    "chimp",
+    "chopsticks",
+    "cockroach",
+    "coffee-mug",
+    "coffin",
+    "coin",
+    "comet",
+    "computer-keyboard",
+    "computer-monitor",
+    "computer-mouse",
+    "conch",
+    "cormorant",
+    "covered-wagon",
+    "cowboy-hat",
+    "crab",
+    "desk-globe",
+    "diamond-ring",
+    "dice",
+    "dog",
+    "dolphin",
+    "doorknob",
+    "drinking-straw",
+    "duck",
+    "dumb-bell",
+    "eiffel-tower",
+    "electric-guitar",
+    "elephant",
+    "elk",
+    "ewer",
+    "eyeglasses",
+    "fern",
+    "fighter-jet",
+    "fire-extinguisher",
+    "fire-hydrant",
+    "fire-truck",
+    "fireworks",
+    "flashlight",
+    "floppy-disk",
+    "football-helmet",
+    "french-horn",
+    "fried-egg",
+    "frisbee",
+    "frog",
+    "frying-pan",
+    "galaxy",
+    "gas-pump",
+    "giraffe",
+    "goat",
+    "golden-gate-bridge",
+    "goldfish",
+    "golf-ball",
+    "goose",
+    "gorilla",
+    "grand-piano",
+    "grapes",
+    "grasshopper",
+    "guitar-pick",
+    "hamburger",
+    "hammock",
+    "harmonica",
+    "harp",
+    "harpsichord",
+    "hawksbill",
+    "head-phones",
+    "helicopter",
+    "hibiscus",
+    "homer-simpson",
+    "horse",
+    "horseshoe-crab",
+    "hot-air-balloon",
+    "hot-dog",
+    "hot-tub",
+    "hourglass",
+    "house-fly",
+    "human-skeleton",
+    "hummingbird",
+    "ibis",
+    "ice-cream-cone",
+    "iguana",
+    "ipod",
+    "iris",
+    "jesus-christ",
+    "joy-stick",
+    "kangaroo",
+    "kayak",
+    "ketch",
+    "killer-whale",
+    "knife",
+    "ladder",
+    "laptop",
+    "lathe",
+    "leopards",
+    "license-plate",
+    "lightbulb",
+    "light-house",
+    "lightning",
+    "llama",
+    "mailbox",
+    "mandolin",
+    "mars",
+    "mattress",
+    "megaphone",
+    "menorah",
+    "microscope",
+    "microwave",
+    "minaret",
+    "minotaur",
+    "motorbikes",
+    "mountain-bike",
+    "mushroom",
+    "mussels",
+    "necktie",
+    "octopus",
+    "ostrich",
+    "owl",
+    "palm-pilot",
+    "palm-tree",
+    "paperclip",
+    "paper-shredder",
+    "pci-card",
+    "penguin",
+    "people",
+    "pez-dispenser",
+    "photocopier",
+    "picnic-table",
+    "playing-card",
+    "porcupine",
+    "pram",
+    "praying-mantis",
+    "pyramid",
+    "raccoon",
+    "radio-telescope",
+    "rainbow",
+    "refrigerator",
+    "revolver",
+    "rifle",
+    "rotary-phone",
+    "roulette-wheel",
+    "saddle",
+    "saturn",
+    "school-bus",
+    "scorpion",
+    "screwdriver",
+    "segway",
+    "self-propelled-lawn-mower",
+    "sextant",
+    "sheet-music",
+    "skateboard",
+    "skunk",
+    "skyscraper",
+    "smokestack",
+    "snail",
+    "snake",
+    "sneaker",
+    "snowmobile",
+    "soccer-ball",
+    "socks",
+    "soda-can",
+    "spaghetti",
+    "speed-boat",
+    "spider",
+    "spoon",
+    "stained-glass",
+    "starfish",
+    "steering-wheel",
+    "stirrups",
+    "sunflower",
+    "superman",
+    "sushi",
+    "swan",
+    "swiss-army-knife",
+    "sword",
+    "syringe",
+    "tambourine",
+    "teapot",
+    "teddy-bear",
+    "teepee",
+    "telephone-box",
+    "tennis-ball",
+    "tennis-court",
+    "tennis-racket",
+    "theodolite",
+    "toaster",
+    "tomato",
+    "tombstone",
+    "top-hat",
+    "touring-bike",
+    "tower-pisa",
+    "traffic-light",
+    "treadmill",
+    "triceratops",
+    "tricycle",
+    "trilobite",
+    "tripod",
+    "t-shirt",
+    "tuning-fork",
+    "tweezer",
+    "umbrella",
+    "unicorn",
+    "vcr",
+    "video-projector",
+    "washing-machine",
+    "watch",
+    "waterfall",
+    "watermelon",
+    "welding-mask",
+    "wheelbarrow",
+    "windmill",
+    "wine-bottle",
+    "xylophone",
+    "yarmulke",
+    "yo-yo",
+    "zebra",
+    "airplanes",
+    "car-side",
+    "faces-easy",
+    "greyhound",
+    "tennis-shoes",
+    "toad",
+]
+imdb = ["Negative", "Positive"]
+mnist = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
+
+urls = {
+    "caltech256": [
+        "https://github.com/cleanlab/label-errors/raw/5392f6c71473055060be3044becdde1cbc18284d/"
+        "original_test_labels/caltech256_original_labels.npy",
+        "https://github.com/cleanlab/label-errors/raw/5392f6c71473055060be3044becdde1cbc18284d"
+        "/cross_validated_predicted_probabilities/caltech256_pyx.npy",
+    ],
+    "mnist": [
+        "https://github.com/cleanlab/label-errors/raw/5392f6c71473055060be3044becdde1cbc18284d"
+        "/original_test_labels/mnist_test_set_original_labels.npy",
+        "https://github.com/cleanlab/label-errors/raw/5392f6c71473055060be3044becdde1cbc18284d"
+        "/cross_validated_predicted_probabilities/mnist_test_set_pyx.npy",
+    ],
+    "imdb": [
+        "https://github.com/cleanlab/label-errors/raw"
+        "/5392f6c71473055060be3044becdde1cbc18284d/original_test_labels"
+        "/imdb_test_set_original_labels.npy",
+        "https://github.com/cleanlab/label-errors/raw/5392f6c71473055060be3044becdde1cbc18284d"
+        "/cross_validated_predicted_probabilities/imdb_test_set_pyx.npy",
+    ],
+    "cifar100": [
+        "https://github.com/cleanlab/label-errors/raw/5392f6c71473055060be3044becdde1cbc18284d"
+        "/original_test_labels/cifar100_test_set_original_labels.npy",
+        "https://github.com/cleanlab/label-errors/raw/5392f6c71473055060be3044becdde1cbc18284d"
+        "/cross_validated_predicted_probabilities/cifar100_test_set_pyx.npy",
+    ],
+}
+
+
+def _get_pred_probs_labels_from_labelerrors_datasets(dataset_name):
+    """Helper function to load data from the labelerrors.com datasets."""
+
+    labels_url, pred_probs_url = urls[dataset_name]
+    response = requests.get(pred_probs_url)
+    response.raise_for_status()
+    pred_probs = np.load(io.BytesIO(response.content), allow_pickle=True)
+    response = requests.get(labels_url)
+    response.raise_for_status()
+    labels = np.load(io.BytesIO(response.content), allow_pickle=True)
+    return pred_probs, labels
+
+
+@pytest.mark.parametrize("dataset_name", ["mnist", "caltech256", "cifar100"])
+def test_real_datasets(dataset_name):
+    print("\n" + dataset_name.capitalize() + "\n")
+    class_names = eval(dataset_name)
+    pred_probs, labels = _get_pred_probs_labels_from_labelerrors_datasets(dataset_name)
+    # if this runs without issue no all four datasets, the test passes
+    _ = health_summary(
+        pred_probs=pred_probs,
+        labels=labels,
+        class_names=class_names,
+        verbose=dataset_name != "mnist",  # test out verbose=False on one of the datasets.
+    )
+
+
+@pytest.mark.parametrize("dataset_name", ["mnist"])
+def test_multilabel_error(dataset_name):
+    print("\n" + dataset_name.capitalize() + "\n")
+    class_names = eval(dataset_name)
+    pred_probs, labels = _get_pred_probs_labels_from_labelerrors_datasets(dataset_name)
+    # if this runs without issue no all four datasets, the test passes
+    with pytest.raises(ValueError) as e:
+        _ = find_overlapping_classes(labels=labels, pred_probs=pred_probs, multi_label=True)
+
+
+@pytest.mark.parametrize("asymmetric", [True, False])
+@pytest.mark.parametrize("dataset_name", ["mnist", "imdb"])
+def test_symmetry_df_size(asymmetric, dataset_name):
+    pred_probs, labels = _get_pred_probs_labels_from_labelerrors_datasets(dataset_name)
+    joint = estimate_joint(labels=labels, pred_probs=pred_probs)
+    num_classes = pred_probs.shape[1]
+    df = find_overlapping_classes(
+        joint=joint,
+        asymmetric=asymmetric,
+        class_names=eval(dataset_name),
+        num_examples=len(labels),
+    )
+    if asymmetric:
+        assert len(df) == num_classes**2 - num_classes
+    else:  # symmetric
+        assert len(df) == (num_classes**2 - num_classes) / 2
+
+        # Second test for symmetric
+        # check that the row, col value returned is actually the sum from the joint.
+        sum_0_1 = joint[0, 1] + joint[1, 0]
+        df_0_1 = df[(df["Class Index A"] == 0) & (df["Class Index B"] == 1)]["Joint Probability"]
+        assert sum_0_1 - df_0_1.values[0] < 1e-8  # Check two floats are equal
+
+
+@pytest.mark.parametrize("use_num_examples", [True, False])
+@pytest.mark.parametrize("use_labels", [True, False])
+@pytest.mark.parametrize(
+    "func", [find_overlapping_classes, rank_classes_by_label_quality, overall_label_health_score]
+)
+def test_value_error_missing_num_examples_with_joint(use_num_examples, use_labels, func):
+    dataset_name = "imdb"
+    pred_probs, labels = _get_pred_probs_labels_from_labelerrors_datasets(dataset_name)
+    joint = estimate_joint(labels=labels, pred_probs=pred_probs)
+    if use_num_examples is False and use_labels is False:  # can't infer num_examples. Throw error!
+        with pytest.raises(ValueError) as e:
+            df = func(
+                labels=labels if use_labels else None,
+                joint=joint,
+                num_examples=len(labels) if use_num_examples else None,
+            )
+    else:  # at least one of use_num_examples and use_labels must be True. Can infer num_examples.
+        # If this runs without error, the test passes.
+        df = func(
+            labels=labels if use_labels else None,
+            joint=joint,
+            num_examples=len(labels) if use_num_examples else None,
+        )
+
+
+@pytest.mark.parametrize("dataset_name", ["mnist", "caltech256", "cifar100"])
+def test_overall_label_health_score_matched_num_issues(dataset_name):
+    # Matches num_label_issues
+    pred_probs, labels = _get_pred_probs_labels_from_labelerrors_datasets(dataset_name)
+    num_issues = num_label_issues(labels=labels, pred_probs=pred_probs)
+    score = overall_label_health_score(labels=labels, pred_probs=pred_probs)
+    assert 1 - num_issues / labels.shape[0] == score
+
+
+def test_overall_label_health_score_function_calls():
+    dataset_name = "caltech256"
+    pred_probs, labels = _get_pred_probs_labels_from_labelerrors_datasets(dataset_name)
+    score = overall_label_health_score(labels=labels, pred_probs=pred_probs)
+
+    confident_joint = compute_confident_joint(labels=labels, pred_probs=pred_probs)
+    num_examples = len(labels)
+    score_cj = overall_label_health_score(
+        labels=None, pred_probs=pred_probs, confident_joint=confident_joint
+    )
+    joint = estimate_joint(labels=labels, pred_probs=pred_probs)
+    score_joint = overall_label_health_score(
+        labels=None, pred_probs=pred_probs, joint=joint, num_examples=num_examples
+    )
+    joint_cj = estimate_joint(labels=labels, pred_probs=pred_probs, confident_joint=confident_joint)
+    score_joint_cj = overall_label_health_score(
+        labels=None, pred_probs=pred_probs, joint=joint_cj, num_examples=num_examples
+    )
+    assert score_cj != score
+    assert score_cj == score_joint
+    assert score_joint_cj == score_joint
+
+
+confident_joint_strategy = npst.arrays(
+    np.int32,
+    shape=npst.array_shapes(min_dims=2, max_dims=2, min_side=2, max_side=10),
+    elements=st.integers(min_value=0, max_value=int(1e6)),
+).filter(lambda arr: arr.shape[0] == arr.shape[1])
+
+
+@pytest.mark.issue_651
+@given(confident_joint=confident_joint_strategy)
+@settings(deadline=500)
+def test_find_overlapping_classes_with_confident_joint(confident_joint):
+    # Setup
+    K = confident_joint.shape[0]
+    overlapping_classes = find_overlapping_classes(confident_joint=confident_joint)
+
+    # Test that the output dataframe has the expected columns
+    expected_columns = [
+        "Class Index A",
+        "Class Index B",
+        "Num Overlapping Examples",
+        "Joint Probability",
+    ]
+    assert set(overlapping_classes.columns) == set(expected_columns)
+
+    # Class indices must be valid
+    assert overlapping_classes["Class Index A"].between(0, K - 1).all()
+    assert overlapping_classes["Class Index B"].between(0, K - 1).all()
+
+    # Overlapping example count should be non-negative integers
+    assert (overlapping_classes["Num Overlapping Examples"] >= 0).all()
+    assert overlapping_classes["Num Overlapping Examples"].dtype == int
+
+    # Joint probabilities should be between 0 and 1
+    assert (overlapping_classes["Joint Probability"] >= 0).all()
+    assert (overlapping_classes["Joint Probability"] <= 1).all()
+
+    # Joint probabilities sorted in descending order
+    if K > 2:
+        assert (overlapping_classes["Joint Probability"].diff()[1:] <= 0).all()

tests/test_frameworks.py

@@ -0,0 +1,143 @@
+"""
+Scripts to test cleanlab usage with various ML frameworks:
+pytorch, skorch
+"""
+
+import pytest
+import warnings
+
+# pytest.mark.filterwarnings is unable to catch filterbuffers library DeprecationWarning
+warnings.filterwarnings(action="ignore", category=DeprecationWarning)
+
+import os
+from copy import deepcopy
+import random
+import numpy as np
+
+import torch
+import skorch
+import sklearn
+
+from cleanlab.classification import CleanLearning
+
+# Version check for sklearn compatibility
+uses_sklearn_ge_1_6_0 = tuple(map(int, sklearn.__version__.split(".")[:2])) >= (1, 6)
+
+
+def dataset_w_errors():
+    num_classes = 2
+    num_features = 3
+    n = 50
+    margin = 5
+    X = np.vstack(
+        [np.random.random((n, num_features)), np.random.random((n, num_features)) + margin]
+    )
+    X = (X - X.mean(axis=0)) / X.std(axis=0)  # normalize columns
+    y = np.array([0] * n + [1] * n)
+    y_og = np.array(y)
+    # Introduce label errors
+    error_indices = [n - 3, n - 2, n - 1, n, n + 1, n + 2]
+    for idx in error_indices:
+        y[idx] = 1 - y[idx]  # Flip label
+
+    if os.name == "nt":  # running on Windows
+        # numpy converts to int32 instead of int64 on Windows, incompatible with neural nets
+        # https://github.com/numpy/numpy/issues/17640
+        X = np.float64(X)
+        y = np.int64(y)
+        y_og = np.int64(y_og)
+
+    return {
+        "X": X,
+        "y": y,
+        "y_og": y_og,
+        "error_indices": error_indices,
+        "num_classes": num_classes,
+        "num_features": num_features,
+    }
+
+
+def make_rare_label(data):
+    """Makes one label really rare in the dataset."""
+    data = deepcopy(data)
+    y = data["y"]
+    class0_inds = np.where(y == 0)[0]
+    if len(class0_inds) < 1:
+        raise ValueError("Class 0 too rare already")
+    class0_inds_remove = class0_inds[1:]
+    if len(class0_inds_remove) > 0:
+        y[class0_inds_remove] = 1
+    data["y"] = y
+    return data
+
+
+SEED = 1
+np.random.seed(SEED)
+random.seed(SEED)
+torch.manual_seed(SEED)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+torch.cuda.manual_seed_all(SEED)
+
+DATA = dataset_w_errors()
+DATA_RARE_LABEL = make_rare_label(DATA)
+
+
+def test_torch(data=DATA, hidden_units=128):
+    dataset = torch.utils.data.TensorDataset(
+        torch.from_numpy(data["X"]).float(), torch.from_numpy(data["y"])
+    )
+
+    class TorchNetwork(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layers = torch.nn.Sequential(
+                torch.nn.Linear(data["num_features"], hidden_units),
+                torch.nn.ReLU(),
+                torch.nn.Linear(hidden_units, data["num_classes"]),
+            )
+
+        def forward(self, X):
+            return self.layers(X)
+
+    # Test base model works:
+    skorch_config = {"criterion": torch.nn.CrossEntropyLoss, "optimizer": torch.optim.Adam}
+    net = skorch.NeuralNet(TorchNetwork, **skorch_config)
+    net.fit(dataset, data["y"], epochs=2)
+    preds_base = net.predict(dataset)
+
+    # Test Cleanlearning performs well:
+    net = skorch.NeuralNet(TorchNetwork, **skorch_config)
+    cl = CleanLearning(net)
+    cl.fit(dataset, data["y"], clf_kwargs={"epochs": 30}, clf_final_kwargs={"epochs": 60})
+
+    preds = cl.predict(dataset).argmax(axis=1)
+    err = np.sum(preds != data["y_og"]) / len(data["y_og"])
+    issue_indices = list(cl.label_issues_df[cl.label_issues_df["is_label_issue"]].index.values)
+    assert issue_indices == data["error_indices"]
+    assert err < 1e-3
+
+
+@pytest.mark.filterwarnings("ignore")
+def test_torch_rarelabel(data=DATA_RARE_LABEL, hidden_units=8):
+    dataset = torch.utils.data.TensorDataset(
+        torch.from_numpy(data["X"]).float(), torch.from_numpy(data["y"])
+    )
+
+    class TorchNetwork(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layers = torch.nn.Sequential(
+                torch.nn.Linear(data["num_features"], hidden_units),
+                torch.nn.ReLU(),
+                torch.nn.Linear(hidden_units, data["num_classes"]),
+            )
+
+        def forward(self, X):
+            return self.layers(X)
+
+    # Test Cleanlearning works:
+    net = skorch.NeuralNet(TorchNetwork, criterion=torch.nn.CrossEntropyLoss)
+    cl = CleanLearning(net)
+    cl.fit(dataset, data["y"], clf_kwargs={"epochs": 2})
+    pred_probs = cl.predict(dataset)

tests/test_latent_algebra.py

@@ -0,0 +1,125 @@
+import numpy as np
+import pytest
+
+from cleanlab.internal import latent_algebra
+
+s = [0] * 10 + [1] * 5 + [2] * 15
+nm = np.array([[1.0, 0.0, 0.2], [0.0, 0.7, 0.2], [0.0, 0.3, 0.6]])
+
+
+def test_latent_py_ps_inv():
+    ps, py, inv = latent_algebra.compute_ps_py_inv_noise_matrix(s, nm)
+    assert all(abs(np.dot(inv, ps) - py) < 1e-3)
+    assert all(abs(np.dot(nm, py) - ps) < 1e-3)
+
+
+def get_latent_py_ps_inv():
+    ps, py, inv = latent_algebra.compute_ps_py_inv_noise_matrix(s, nm)
+
+    return ps, py, inv
+
+
+def test_latent_inv():
+    ps, py, inv = get_latent_py_ps_inv()
+    inv2 = latent_algebra.compute_inv_noise_matrix(py, nm)
+    assert np.all(abs(inv - inv2) < 1e-3)
+
+
+def test_latent_nm():
+    ps, py, inv = get_latent_py_ps_inv()
+    nm2 = latent_algebra.compute_noise_matrix_from_inverse(ps, inv, py=py)
+    assert np.all(abs(nm - nm2) < 1e-3)
+
+
+def test_latent_py():
+    ps, py, inv = get_latent_py_ps_inv()
+    py2 = latent_algebra.compute_py(ps, nm, inv)
+    assert np.all(abs(py - py2) < 1e-3)
+
+
+def test_latent_py_warning():
+    ps, py, inv = get_latent_py_ps_inv()
+    with pytest.raises(TypeError) as e:
+        with pytest.warns(UserWarning) as w:
+            py2 = latent_algebra.compute_py(
+                ps=np.array([[[0.1, 0.3, 0.6]]]),
+                noise_matrix=nm,
+                inverse_noise_matrix=inv,
+            )
+            py2 = latent_algebra.compute_py(
+                ps=np.array([[0.1], [0.2], [0.7]]),
+                noise_matrix=nm,
+                inverse_noise_matrix=inv,
+            )
+
+
+def test_compute_py_err():
+    ps, py, inv = get_latent_py_ps_inv()
+    try:
+        py = latent_algebra.compute_py(
+            ps=ps,
+            noise_matrix=nm,
+            inverse_noise_matrix=inv,
+            py_method="marginal_ps",
+        )
+    except ValueError as e:
+        assert "true_labels_class_counts" in str(e)
+        with pytest.raises(ValueError) as e:
+            py = latent_algebra.compute_py(
+                ps=ps,
+                noise_matrix=nm,
+                inverse_noise_matrix=inv,
+                py_method="marginal_ps",
+            )
+
+
+def test_compute_py_marginal_ps():
+    ps, py, inv = get_latent_py_ps_inv()
+    cj = nm * ps * len(s)
+    true_labels_class_counts = cj.sum(axis=0)
+    py2 = latent_algebra.compute_py(
+        ps=ps,
+        noise_matrix=nm,
+        inverse_noise_matrix=inv,
+        py_method="marginal_ps",
+        true_labels_class_counts=true_labels_class_counts,
+    )
+    assert all(abs(py - py2) < 1e-2)
+
+
+def test_pyx():
+    pred_probs = np.array(
+        [
+            [0.1, 0.3, 0.6],
+            [0.1, 0.0, 0.9],
+            [0.1, 0.0, 0.9],
+            [1.0, 0.0, 0.0],
+            [0.1, 0.8, 0.1],
+        ]
+    )
+    ps, py, inv = get_latent_py_ps_inv()
+    pyx = latent_algebra.compute_pyx(pred_probs, nm, inv)
+    assert np.all(np.sum(pyx, axis=1) - 1 < 1e-4)
+
+
+def test_pyx_error():
+    pred_probs = np.array([0.1, 0.3, 0.6])
+    ps, py, inv = get_latent_py_ps_inv()
+    try:
+        pyx = latent_algebra.compute_pyx(pred_probs, nm, inverse_noise_matrix=inv)
+    except ValueError as e:
+        assert "should be (N, K)" in str(e)
+    with pytest.raises(ValueError) as e:
+        pyx = latent_algebra.compute_pyx(pred_probs, nm, inverse_noise_matrix=inv)
+
+
+def test_compute_py_method_marginal_true_labels_class_counts_none_error():
+    ps, py, inv = get_latent_py_ps_inv()
+    with pytest.raises(ValueError) as e:
+        _ = latent_algebra.compute_py(
+            ps=ps,
+            noise_matrix=nm,
+            inverse_noise_matrix=inv,
+            py_method="marginal",
+            true_labels_class_counts=None,
+        )

tests/test_model_pytorch_cnn.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+from cleanlab.experimental.mnist_pytorch import (
+    CNN,
+    SKLEARN_DIGITS_TEST_SIZE,
+    SKLEARN_DIGITS_TRAIN_SIZE,
+)
+import cleanlab
+import numpy as np
+from sklearn.metrics import accuracy_score
+from sklearn.datasets import load_digits
+import pytest
+
+X_train_idx = np.arange(SKLEARN_DIGITS_TRAIN_SIZE)
+X_test_idx = np.arange(SKLEARN_DIGITS_TEST_SIZE)
+# Get sklearn digits data labels
+_, y_all = load_digits(return_X_y=True)
+# PyTorch requires type long targets.
+y_train = y_all[:-SKLEARN_DIGITS_TEST_SIZE].astype(np.int32)
+true_labels_test = y_all[-SKLEARN_DIGITS_TEST_SIZE:].astype(np.int32)
+
+
+@pytest.mark.slow
+def test_loaders(
+    seed=0,
+):
+    """This is going to OVERFIT - train and test on the SAME SET.
+    The goal of this test is just to make sure the data loads correctly.
+    And all the main functions work."""
+
+    from cleanlab.count import estimate_confident_joint_and_cv_pred_proba, estimate_latent
+
+    np.random.seed(seed)
+    filter_by = "prune_by_noise_rate"
+    # Pre-train for only 3 epochs (it maxes out around 8-12 epochs)
+    cnn = CNN(epochs=3, log_interval=None, seed=seed, dataset="sklearn-digits")
+    score = 0
+    for loader in ["train", "test", None]:
+        print("loader:", loader)
+        prev_score = score
+        X = X_test_idx if loader == "test" else X_train_idx
+        y = true_labels_test if loader == "test" else y_train
+        # Setting this overrides all future functions.
+        cnn.loader = loader
+        # pre-train (overfit, not out-of-sample) to entire dataset.
+        cnn.fit(
+            X,
+            None,
+        )
+        # This next block of code checks if cleanlab works with the CNN
+        # Out-of-sample cross-validated holdout predicted probabilities
+        np.random.seed(seed)
+        # Single epoch for cross-validation (already pre-trained)
+        cnn.epochs = 1
+        cj, pred_probs = estimate_confident_joint_and_cv_pred_proba(X, y, cnn, cv_n_folds=2)
+        est_py, est_nm, est_inv = estimate_latent(cj, y)
+        # algorithmic identification of label issues
+        err_idx = cleanlab.filter.find_label_issues(
+            y, pred_probs, confident_joint=cj, filter_by=filter_by
+        )
+        assert err_idx is not None
+
+        # Get prediction on loader set.
+        pred = cnn.predict(X)
+        score = accuracy_score(y, pred)
+        print("Acc Before: {:.2f}, After: {:.2f}".format(prev_score, score))
+        assert score > prev_score  # Scores should increase
+
+
+def test_throw_exception():
+    cnn = CNN(epochs=1, log_interval=1000, seed=0)
+    try:
+        cnn.fit(train_idx=[0, 1], train_labels=[1])
+    except Exception as e:
+        assert "same length" in str(e)
+        with pytest.raises(ValueError) as e:
+            cnn.fit(train_idx=[0, 1], train_labels=[1])
+
+
+def test_n_train_examples():
+    cnn = CNN(
+        epochs=4,
+        log_interval=1000,
+        loader="train",
+        seed=1,
+        dataset="sklearn-digits",
+    )
+    cnn.fit(
+        train_idx=X_train_idx,
+        train_labels=y_train,
+        loader="train",
+    )
+    cnn.loader = "test"
+    pred = cnn.predict(X_test_idx)
+    print(accuracy_score(true_labels_test, pred))
+    assert accuracy_score(true_labels_test, pred) > 0.1
+
+    # Check that exception is raised when invalid name is given.
+    cnn.loader = "INVALID"
+    with pytest.raises(ValueError) as e:
+        pred = cnn.predict(X_test_idx)
+
+    # Check that pred_proba runs on all examples when None is passed in
+    cnn.loader = "test"
+    proba = cnn.predict_proba(idx=None, loader="test")
+    assert proba is not None
+    assert len(pred) == SKLEARN_DIGITS_TEST_SIZE

tests/test_multiannotator.py

@@ -0,0 +1,822 @@
+from copy import deepcopy
+
+import numpy as np
+import pandas as pd
+import pytest
+from sklearn.linear_model import LogisticRegression
+
+from cleanlab import count
+from cleanlab.benchmarking.noise_generation import (
+    generate_noise_matrix_from_trace,
+    generate_noisy_labels,
+)
+from cleanlab.internal.multiannotator_utils import (
+    assert_valid_inputs_multiannotator,
+    format_multiannotator_labels,
+)
+from cleanlab.multiannotator import (
+    convert_long_to_wide_dataset,
+    get_active_learning_scores,
+    get_active_learning_scores_ensemble,
+    get_label_quality_multiannotator,
+    get_label_quality_multiannotator_ensemble,
+    get_majority_vote_label,
+)
+
+
+def make_data(
+    means=[[3, 2], [7, 7], [0, 8]],
+    covs=[[[5, -1.5], [-1.5, 1]], [[1, 0.5], [0.5, 4]], [[5, 1], [1, 5]]],
+    labeled_sizes=[80, 40, 40],
+    unlabeled_sizes=[20, 10, 10],
+    avg_trace=0.8,
+    num_annotators=50,
+    seed=1,  # set to None for non-reproducible randomness
+):
+    np.random.seed(seed=seed)
+
+    m = len(means)  # number of classes
+    n = sum(labeled_sizes)
+    local_data = []
+    labels = []
+    unlabeled_data = []
+    unlabeled_labels = []
+
+    for idx in range(m):
+        local_data.append(
+            np.random.multivariate_normal(mean=means[idx], cov=covs[idx], size=labeled_sizes[idx])
+        )
+        unlabeled_data.append(
+            np.random.multivariate_normal(mean=means[idx], cov=covs[idx], size=unlabeled_sizes[idx])
+        )
+        labels.append(np.array([idx for i in range(labeled_sizes[idx])]))
+        unlabeled_labels.append(np.array([idx for i in range(unlabeled_sizes[idx])]))
+    X_train = np.vstack(local_data)
+    X_train_unlabeled = np.vstack(unlabeled_data)
+    true_labels_train = np.hstack(labels)
+    true_labels_train_unlabeled = np.hstack(unlabeled_labels)
+
+    # Compute p(true_label=k)
+    py = np.bincount(true_labels_train) / float(len(true_labels_train))
+
+    noise_matrix = generate_noise_matrix_from_trace(
+        m,
+        trace=avg_trace * m,
+        py=py,
+        valid_noise_matrix=True,
+        seed=seed,
+    )
+
+    # Generate our noisy labels using the noise_matrix for specified number of annotators.
+    s = pd.DataFrame(
+        np.vstack(
+            [generate_noisy_labels(true_labels_train, noise_matrix) for _ in range(num_annotators)]
+        ).transpose()
+    )
+
+    # column of labels without NaNs to test _get_worst_class
+    complete_labels = deepcopy(s)
+
+    # Each annotator only labels approximately 20% of the dataset
+    # (unlabeled points represented with NaN)
+    s = s.apply(lambda x: x.mask(np.random.random(n) < 0.8))
+    s.dropna(axis=1, how="all", inplace=True)
+
+    # Estimate pred_probs
+    latent = count.estimate_py_noise_matrices_and_cv_pred_proba(
+        X=X_train,
+        labels=true_labels_train,
+        cv_n_folds=3,
+    )
+    latent_unlabeled = count.estimate_py_noise_matrices_and_cv_pred_proba(
+        X=X_train_unlabeled,
+        labels=true_labels_train_unlabeled,
+        cv_n_folds=3,
+    )
+
+    row_NA_check = pd.notna(s).any(axis=1)
+
+    return {
+        "X_train": X_train[row_NA_check],
+        "X_train_unlabeled": X_train_unlabeled,
+        "X_train_complete": X_train,
+        "true_labels_train": true_labels_train[row_NA_check],
+        "true_labels_train_unlabeled": true_labels_train_unlabeled,
+        "labels": s[row_NA_check].reset_index(drop=True),
+        "labels_unlabeled": pd.DataFrame(
+            np.full((len(true_labels_train_unlabeled), num_annotators), np.nan)
+        ),
+        "complete_labels": complete_labels,
+        "pred_probs": latent[4][row_NA_check],
+        "pred_probs_unlabeled": latent_unlabeled[4],
+        "pred_probs_complete": latent[4],
+        "noise_matrix": noise_matrix,
+    }
+
+
+def make_ensemble_data(
+    means=[[3, 2], [7, 7], [0, 8]],
+    covs=[[[5, -1.5], [-1.5, 1]], [[1, 0.5], [0.5, 4]], [[5, 1], [1, 5]]],
+    unlabeled_sizes=[20, 10, 10],
+    avg_trace=0.8,
+    num_annotators=50,
+    seed=1,  # set to None for non-reproducible randomness
+):
+    np.random.seed(seed=seed)
+
+    data = make_data()
+
+    X_train = data["X_train"]
+    true_labels_train = data["true_labels_train"]
+    X_train_unlabeled = data["X_train_unlabeled"]
+    true_labels_train_unlabeled = data["true_labels_train_unlabeled"]
+
+    # Estimate pred_probs for unlabeled data
+    pred_probs_extra = count.estimate_py_noise_matrices_and_cv_pred_proba(
+        X=X_train,
+        labels=true_labels_train,
+        cv_n_folds=3,
+        clf=LogisticRegression(),
+    )[4]
+    pred_probs_labeled = np.array([data["pred_probs"], pred_probs_extra])
+
+    # Estimate pred_probs for labeled data
+    pred_probs_extra_unlabeled = count.estimate_py_noise_matrices_and_cv_pred_proba(
+        X=X_train_unlabeled,
+        labels=true_labels_train_unlabeled,
+        cv_n_folds=3,
+        clf=LogisticRegression(),
+    )[4]
+    pred_probs_unlabeled = np.array([data["pred_probs_unlabeled"], pred_probs_extra_unlabeled])
+
+    return {
+        "X_train": data["X_train"],
+        "X_train_unlabeled": data["X_train_unlabeled"],
+        "true_labels_train": data["true_labels_train"],
+        "true_labels_train_unlabeled": data["true_labels_train_unlabeled"],
+        "labels": data["labels"],
+        "labels_unlabeled": data["labels_unlabeled"],
+        "complete_labels": data["complete_labels"],
+        "pred_probs": pred_probs_labeled,
+        "pred_probs_unlabeled": pred_probs_unlabeled,
+        "noise_matrix": data["noise_matrix"],
+    }
+
+
+def make_data_long(data):
+    data_long = data.stack().reset_index()
+    data_long.columns = ["task", "annotator", "label"]
+
+    return data_long
+
+
+# Global to be used by all test methods. Only compute this once for speed.
+data = make_data()
+ensemble_data = make_ensemble_data()
+small_data = make_data(
+    labeled_sizes=[5, 5, 5],
+    unlabeled_sizes=[5, 5, 5],
+    num_annotators=1,
+)
+
+
+def test_convert_long_to_wide():
+    labels_long = make_data_long(data["labels"])
+    labels_wide = convert_long_to_wide_dataset(labels_long)
+
+    assert isinstance(labels_wide, pd.DataFrame)
+
+    # ensures labels_long contains all the non-NaN values of labels_wide
+    # Account for different pandas stack() behaviors: 2.x drops NaN, 3.x preserves NaN
+    expected_count = labels_long["label"].notna().sum()
+    assert labels_wide.count(axis=1).sum() == expected_count
+
+    # checks one index to make sure data is consistent across both dataframes
+    example_long = labels_long[labels_long["task"] == 0]
+    # Only compare non-null entries to handle both pandas behaviors
+    example_long_non_null = example_long[example_long["label"].notna()].sort_values("annotator")
+    example_wide = labels_wide.iloc[0].dropna()
+    assert all(example_long_non_null["annotator"] == example_wide.index)
+    assert all(
+        example_long_non_null["label"].reset_index(drop=True) == example_wide.reset_index(drop=True)
+    )
+
+
+def test_label_quality_scores_multiannotator():
+    labels = data["labels"]
+    pred_probs = data["pred_probs"]
+
+    multiannotator_dict = get_label_quality_multiannotator(labels, pred_probs)
+    assert isinstance(multiannotator_dict, dict)
+    assert len(multiannotator_dict) == 3
+
+    label_quality_multiannotator = multiannotator_dict["label_quality"]
+    assert isinstance(label_quality_multiannotator, pd.DataFrame)
+    assert len(label_quality_multiannotator) == len(labels)
+    assert all(label_quality_multiannotator["num_annotations"] > 0)
+    assert set(label_quality_multiannotator["consensus_label"]).issubset(np.unique(labels))
+    assert all(
+        (label_quality_multiannotator["annotator_agreement"] >= 0)
+        & (label_quality_multiannotator["annotator_agreement"] <= 1)
+    )
+    assert all(
+        (label_quality_multiannotator["consensus_quality_score"] >= 0)
+        & (label_quality_multiannotator["consensus_quality_score"] <= 1)
+    )
+
+    annotator_stats = multiannotator_dict["annotator_stats"]
+    assert isinstance(annotator_stats, pd.DataFrame)
+    assert len(annotator_stats) == labels.shape[1]
+    assert all(
+        (annotator_stats["annotator_quality"] >= 0) & (annotator_stats["annotator_quality"] <= 1)
+    )
+    assert all(annotator_stats["num_examples_labeled"] > 0)
+    assert all(
+        (annotator_stats["agreement_with_consensus"] >= 0)
+        & (annotator_stats["agreement_with_consensus"] <= 1)
+    )
+    assert set(annotator_stats["worst_class"]).issubset(np.unique(labels))
+
+    detailed_label_quality = multiannotator_dict["detailed_label_quality"]
+    assert detailed_label_quality.shape == labels.shape
+
+    # test verbose=False
+    multiannotator_dict = get_label_quality_multiannotator(labels, pred_probs, verbose=False)
+
+    # test passing a list into consensus_method
+    multiannotator_dict = get_label_quality_multiannotator(
+        labels, pred_probs, consensus_method=["majority_vote", "best_quality"]
+    )
+
+    # test passing arguments for get_label_quality_scores
+    multiannotator_dict = get_label_quality_multiannotator(
+        labels, pred_probs, label_quality_score_kwargs={"method": "normalized_margin"}
+    )
+
+    # test different quality_methods
+    # also testing passing labels as np.ndarray
+    multiannotator_dict = get_label_quality_multiannotator(
+        np.array(labels), pred_probs, quality_method="agreement"
+    )
+
+    # test returning annotator_stats
+    multiannotator_dict = get_label_quality_multiannotator(
+        labels, pred_probs, return_annotator_stats=False
+    )
+    assert isinstance(multiannotator_dict, dict)
+    assert len(multiannotator_dict) == 2
+    assert isinstance(multiannotator_dict["label_quality"], pd.DataFrame)
+    assert isinstance(multiannotator_dict["detailed_label_quality"], pd.DataFrame)
+
+    # test returning detailed_label_quality
+    multiannotator_dict = get_label_quality_multiannotator(
+        labels, pred_probs, return_detailed_quality=False
+    )
+    assert isinstance(multiannotator_dict, dict)
+    assert len(multiannotator_dict) == 2
+    assert isinstance(multiannotator_dict["label_quality"], pd.DataFrame)
+    assert isinstance(multiannotator_dict["annotator_stats"], pd.DataFrame)
+
+    # test return detailed and annotator stats
+    multiannotator_dict = get_label_quality_multiannotator(
+        labels, pred_probs, return_detailed_quality=False, return_annotator_stats=False
+    )
+    assert isinstance(multiannotator_dict, dict)
+    assert len(multiannotator_dict) == 1
+    assert isinstance(multiannotator_dict["label_quality"], pd.DataFrame)
+
+    # test return model and annotator weights
+    multiannotator_dict = get_label_quality_multiannotator(labels, pred_probs, return_weights=True)
+    assert len(multiannotator_dict) == 5
+    assert isinstance(multiannotator_dict["model_weight"], float)
+    assert isinstance(multiannotator_dict["annotator_weight"], np.ndarray)
+
+    # test non-numeric annotator names
+    labels_string_names = labels.add_prefix("anno_")
+    multiannotator_dict = get_label_quality_multiannotator(
+        labels_string_names, pred_probs, return_detailed_quality=False
+    )
+
+    # test calibration
+    multiannotator_dict = get_label_quality_multiannotator(labels, pred_probs, calibrate_probs=True)
+
+    # test incorrect consensus_method
+    try:
+        multiannotator_dict = get_label_quality_multiannotator(
+            labels, pred_probs, consensus_method="fake_method"
+        )
+    except ValueError as e:
+        assert "not a valid consensus method" in str(e)
+
+    # test error when return_weights == True and quality_method != "crowdlab"
+    try:
+        multiannotator_dict = get_label_quality_multiannotator(
+            labels, pred_probs, return_weights=True, quality_method="agreement"
+        )
+    except ValueError as e:
+        assert (
+            "Model and annotator weights are only applicable to the crowdlab quality method"
+            in str(e)
+        )
+
+    # test error catching when labels_multiannotator has NaN columns
+    labels_NA = deepcopy(labels_string_names)
+    labels_NA["anno_0"] = pd.NA
+    try:
+        multiannotator_dict = get_label_quality_multiannotator(
+            labels_NA,
+            pred_probs,
+        )
+    except ValueError as e:
+        assert "cannot have columns with all NaN" in str(e)
+        assert "Annotators ['anno_0'] did not label any examples." in str(e)
+
+    # try same thing as above but with numpy array
+    labels_nan = deepcopy(labels).values.astype(float)
+    labels_nan[:, 1] = np.nan
+    try:
+        multiannotator_dict = get_label_quality_multiannotator(
+            labels_nan,
+            pred_probs,
+        )
+    except ValueError as e:
+        assert "cannot have columns with all NaN" in str(e)
+        assert (
+            "Annotators [" in str(e) and "1" in str(e) and "] did not label any examples." in str(e)
+        )
+
+    # test error catching when labels_multiannotator has NaN rows
+    labels_nan = pd.DataFrame(
+        [
+            [0, np.nan, np.nan],
+            [np.nan, 1, np.nan],
+            [np.nan, np.nan, 2],
+            [np.nan, np.nan, np.nan],
+            [np.nan, np.nan, 2],
+        ]
+    )
+    pred_probs = np.random.random((5, 3))
+
+    try:
+        multiannotator_dict = get_label_quality_multiannotator(
+            labels_nan,
+            pred_probs,
+        )
+    except ValueError as e:
+        assert "cannot have rows with all NaN" in str(e)
+        assert "Examples [" in str(e) and "3" in str(e) and "] do not have any labels." in str(e)
+
+    # test error when using wrong function
+    try:
+        multiannotator_dict = get_label_quality_multiannotator(
+            labels, np.array([pred_probs, pred_probs]), return_weights=True
+        )
+    except ValueError as e:
+        assert "use the ensemble version of this function" in str(e)
+
+    # make sure error is thrown if labels are not 2D
+    labels_flat = labels.values[:, 0].flatten()
+    print(labels_flat.ndim)
+    print(labels_flat)
+    try:
+        multiannotator_dict = get_label_quality_multiannotator(labels_flat, pred_probs)
+    except ValueError as e:
+        assert "labels_multiannotator must be a 2D array or dataframe" in str(e)
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_label_quality_scores_multiannotator_ensemble():
+    labels = ensemble_data["labels"]
+    pred_probs = ensemble_data["pred_probs"]
+
+    multiannotator_dict = get_label_quality_multiannotator_ensemble(
+        labels, pred_probs, return_weights=True
+    )
+    assert isinstance(multiannotator_dict, dict)
+    assert len(multiannotator_dict) == 5
+    assert isinstance(multiannotator_dict["label_quality"], pd.DataFrame)
+    assert isinstance(multiannotator_dict["annotator_stats"], pd.DataFrame)
+    assert isinstance(multiannotator_dict["detailed_label_quality"], pd.DataFrame)
+    assert isinstance(multiannotator_dict["model_weight"], np.ndarray)
+    assert isinstance(multiannotator_dict["annotator_weight"], np.ndarray)
+
+    # test non-numeric annotator names
+    labels_string_names = labels.add_prefix("anno_")
+    multiannotator_dict = get_label_quality_multiannotator_ensemble(
+        labels_string_names, pred_probs, return_detailed_quality=False
+    )
+
+    # test return model and annotator weights
+    multiannotator_dict = get_label_quality_multiannotator_ensemble(
+        labels, pred_probs, return_weights=True
+    )
+    assert len(multiannotator_dict) == 5
+    assert isinstance(multiannotator_dict["model_weight"], np.ndarray)
+    assert isinstance(multiannotator_dict["annotator_weight"], np.ndarray)
+
+    # test numpy arrays and calibrationg
+    multiannotator_dict = get_label_quality_multiannotator_ensemble(
+        np.array(labels), pred_probs, calibrate_probs=True
+    )
+
+    # testing tiebreaks in ensemble
+    labels_tiebreaks = np.array([[1, 2, 0], [1, 1, 0], [1, 0, 0], [2, 2, 2], [1, 2, 0], [1, 2, 0]])
+    pred_probs_tiebreaks = np.array(
+        [
+            [0.4, 0.4, 0.2],
+            [0.3, 0.6, 0.1],
+            [0.75, 0.2, 0.05],
+            [0.1, 0.4, 0.5],
+            [0.2, 0.4, 0.4],
+            [0.2, 0.4, 0.4],
+        ]
+    )
+    pred_probs_tiebreaks_ensemble = np.array(
+        [pred_probs_tiebreaks, pred_probs_tiebreaks, pred_probs_tiebreaks]
+    )
+
+    consensus_label = get_label_quality_multiannotator_ensemble(
+        labels_tiebreaks, pred_probs_tiebreaks_ensemble
+    )
+
+    # test error when using wrong function
+    try:
+        multiannotator_dict = get_label_quality_multiannotator_ensemble(
+            labels, pred_probs[0], return_weights=True
+        )
+    except ValueError as e:
+        assert "use the non-ensemble version of this function" in str(e)
+
+
+def test_get_active_learning_scores():
+    labels = data["labels"]
+    pred_probs = data["pred_probs"]
+    pred_probs_unlabeled = data["pred_probs_unlabeled"]
+
+    # test default case
+    active_learning_scores, active_learning_scores_unlabeled = get_active_learning_scores(
+        labels, pred_probs, pred_probs_unlabeled
+    )
+    assert isinstance(active_learning_scores, np.ndarray)
+    assert len(active_learning_scores) == len(pred_probs)
+    assert len(active_learning_scores_unlabeled) == len(pred_probs_unlabeled)
+
+    # test case where all examples are already labeled
+    # also tests passing labels as np array
+    active_learning_scores, active_learning_scores_unlabeled = get_active_learning_scores(
+        np.array(labels), pred_probs
+    )
+    assert isinstance(active_learning_scores, np.ndarray)
+    assert len(active_learning_scores) == len(pred_probs)
+    assert len(active_learning_scores_unlabeled) == 0
+
+    # test case where only passing unlabeled examples
+    active_learning_scores, active_learning_scores_unlabeled = get_active_learning_scores(
+        pred_probs_unlabeled=pred_probs_unlabeled
+    )
+    assert len(active_learning_scores) == 0
+    assert len(active_learning_scores_unlabeled) == len(pred_probs_unlabeled)
+
+    # test case where number of classes do not match
+    try:
+        active_learning_scores, active_learning_scores_unlabeled = get_active_learning_scores(
+            labels, pred_probs, pred_probs_unlabeled[:, :-1]
+        )
+    except ValueError as e:
+        assert "must have the same number of classes" in str(e)
+
+    # test starting with single labeled example + one unlabeled example
+    single_labels = data["complete_labels"].iloc[[0]]
+    singe_pred_probs = pred_probs[[0]]
+    singe_pred_probs_unlabeled = pred_probs_unlabeled[[0]]
+    get_active_learning_scores(single_labels, singe_pred_probs, singe_pred_probs_unlabeled)
+
+    # test when each example is only labeled by one annotator
+    labels = pd.DataFrame(
+        [
+            [0, np.nan, np.nan],
+            [np.nan, 1, np.nan],
+            [np.nan, np.nan, 2],
+            [np.nan, 1, np.nan],
+            [np.nan, np.nan, 2],
+        ]
+    )
+    pred_probs = np.random.random((5, 3))
+    get_active_learning_scores(labels, pred_probs, pred_probs)
+
+
+def test_get_active_learning_scores_ensemble():
+    labels = ensemble_data["labels"]
+    pred_probs = ensemble_data["pred_probs"]
+    labels_unlabeled = ensemble_data["labels_unlabeled"]
+    pred_probs_unlabeled = ensemble_data["pred_probs_unlabeled"]
+
+    # test default case
+    active_learning_scores, active_learning_scores_unlabeled = get_active_learning_scores_ensemble(
+        labels, pred_probs, pred_probs_unlabeled
+    )
+    assert isinstance(active_learning_scores, np.ndarray)
+    assert len(active_learning_scores) == len(labels)
+    assert len(active_learning_scores_unlabeled) == pred_probs_unlabeled.shape[1]
+
+    # test case where all examples are already labeled
+    # also tests passing labels as np array
+    active_learning_scores, active_learning_scores_unlabeled = get_active_learning_scores_ensemble(
+        np.array(labels), pred_probs
+    )
+    assert isinstance(active_learning_scores, np.ndarray)
+    assert len(active_learning_scores) == len(labels)
+    assert len(active_learning_scores_unlabeled) == 0
+
+    # test case where only passing unlabeled examples
+    active_learning_scores, active_learning_scores_unlabeled = get_active_learning_scores_ensemble(
+        pred_probs_unlabeled=pred_probs_unlabeled
+    )
+    assert len(active_learning_scores) == 0
+    assert len(active_learning_scores_unlabeled) == len(labels_unlabeled)
+
+    # test case where number of classes do not match
+    try:
+        (
+            active_learning_scores,
+            active_learning_scores_unlabeled,
+        ) = get_active_learning_scores_ensemble(labels, pred_probs, pred_probs_unlabeled[:, :-1])
+    except ValueError as e:
+        assert "must have the same number of classes" in str(e)
+
+    # test starting with single labeled example + one unlabeled example
+    single_labels = ensemble_data["complete_labels"].iloc[[0]]
+    singe_pred_probs = pred_probs[:, [0]]
+    singe_pred_probs_unlabeled = pred_probs_unlabeled[:, [0]]
+    get_active_learning_scores_ensemble(single_labels, singe_pred_probs, singe_pred_probs_unlabeled)
+
+    # test when each example is only labeled by one annotator
+    labels = pd.DataFrame(
+        [
+            [0, np.nan, np.nan],
+            [np.nan, 1, np.nan],
+            [np.nan, np.nan, 2],
+            [np.nan, 1, np.nan],
+            [np.nan, np.nan, 2],
+        ]
+    )
+    pred_probs = np.random.random((2, 5, 3))
+
+    get_active_learning_scores_ensemble(labels, pred_probs)
+
+
+def test_single_label_active_learning():
+    labels = np.array(small_data["complete_labels"])
+    labels_unlabeled = small_data["true_labels_train_unlabeled"]
+    pred_probs = small_data["pred_probs_complete"]
+    pred_probs_unlabeled = small_data["pred_probs_unlabeled"]
+
+    assert len(labels) == 15
+
+    # test 5 rounds of active learning
+    for i in range(5):
+        active_learning_scores, active_learning_scores_unlabeled = get_active_learning_scores(
+            labels, pred_probs, pred_probs_unlabeled
+        )
+
+        min_ind = np.argmin(active_learning_scores_unlabeled)
+
+        labels = np.append(labels, labels_unlabeled[min_ind]).reshape(-1, 1)
+        pred_probs = np.append(pred_probs, pred_probs_unlabeled[min_ind].reshape(1, -1), axis=0)
+        labels_unlabeled = np.delete(labels_unlabeled, min_ind)
+        pred_probs_unlabeled = np.delete(pred_probs_unlabeled, min_ind, axis=0)
+
+    assert len(labels) == 20
+
+    # make sure error is thrown if labels are not 2D
+    labels_flat = np.array(small_data["complete_labels"]).reshape(1, -1)
+    try:
+        active_learning_scores, active_learning_scores_unlabeled = get_active_learning_scores(
+            labels, pred_probs, pred_probs_unlabeled
+        )
+    except ValueError as e:
+        assert "labels_multiannotator must be a 2D array or dataframe" in str(e)
+
+
+def test_single_label_active_learning_ensemble():
+    labels = np.array(small_data["complete_labels"])
+    labels_unlabeled = small_data["true_labels_train_unlabeled"]
+    pred_probs = small_data["pred_probs_complete"]
+    pred_probs_unlabeled = small_data["pred_probs_unlabeled"]
+
+    assert len(labels) == 15
+
+    # test 5 rounds of active learning
+    for i in range(5):
+        (
+            active_learning_scores,
+            active_learning_scores_unlabeled,
+        ) = get_active_learning_scores_ensemble(
+            labels,
+            np.array([pred_probs, pred_probs]),
+            np.array([pred_probs_unlabeled, pred_probs_unlabeled]),
+        )
+
+        min_ind = np.argmin(active_learning_scores_unlabeled)
+
+        labels = np.append(labels, labels_unlabeled[min_ind]).reshape(-1, 1)
+        pred_probs = np.append(pred_probs, pred_probs_unlabeled[min_ind].reshape(1, -1), axis=0)
+        labels_unlabeled = np.delete(labels_unlabeled, min_ind)
+        pred_probs_unlabeled = np.delete(pred_probs_unlabeled, min_ind, axis=0)
+
+    assert len(labels) == 20
+
+    # make sure error is thrown if labels are not 2D
+    labels_flat = np.array(small_data["complete_labels"]).reshape(1, -1)
+    try:
+        (
+            active_learning_scores,
+            active_learning_scores_unlabeled,
+        ) = get_active_learning_scores_ensemble(
+            labels,
+            np.array([pred_probs, pred_probs]),
+            np.array([pred_probs_unlabeled, pred_probs_unlabeled]),
+        )
+    except ValueError as e:
+        assert "labels_multiannotator must be a 2D array or dataframe" in str(e)
+
+
+def test_missing_class():
+    labels = np.array(
+        [
+            [1, np.nan, 2],
+            [1, 1, 2],
+            [2, 2, 1],
+            [np.nan, 2, 2],
+            [np.nan, 2, 1],
+            [np.nan, 2, 2],
+        ]
+    )
+
+    pred_probs = np.array(
+        [
+            [0.4, 0.4, 0.2],
+            [0.3, 0.6, 0.1],
+            [0.05, 0.2, 0.75],
+            [0.1, 0.4, 0.5],
+            [0.2, 0.4, 0.4],
+            [0.2, 0.4, 0.4],
+        ]
+    )
+
+    # test default case
+    consensus_label = get_majority_vote_label(labels)
+    consensus_label = get_majority_vote_label(labels, pred_probs)
+    multiannotator_dict = get_label_quality_multiannotator(labels, pred_probs)
+
+    # test other consensus and quality methods
+    multiannotator_dict = get_label_quality_multiannotator(
+        labels, pred_probs, quality_method="agreement"
+    )
+    multiannotator_dict = get_label_quality_multiannotator(
+        labels, pred_probs, consensus_method="majority_vote"
+    )
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_rare_class():
+    labels = np.array(
+        [
+            [1, np.nan, 2],
+            [1, 1, 0],
+            [2, 2, 0],
+            [np.nan, 2, 2],
+            [np.nan, 2, 1],
+            [np.nan, 2, 2],
+        ]
+    )
+
+    pred_probs = np.array(
+        [
+            [0.4, 0.4, 0.2],
+            [0.3, 0.6, 0.1],
+            [0.05, 0.2, 0.75],
+            [0.1, 0.4, 0.5],
+            [0.2, 0.4, 0.4],
+            [0.2, 0.4, 0.4],
+        ]
+    )
+
+    consensus_label = get_majority_vote_label(labels)
+    multiannotator_dict = get_label_quality_multiannotator(labels, pred_probs)
+
+    pred_probs_missing = np.array(
+        [
+            [0.8, 0.2],
+            [0.6, 0.14],
+            [0.95, 0.05],
+            [0.5, 0.5],
+            [0.4, 0.6],
+            [0.4, 0.6],
+        ]
+    )
+
+    try:
+        multiannotator_dict = get_label_quality_multiannotator(labels, pred_probs_missing)
+    except ValueError as e:
+        assert "pred_probs must have at least 3 columns" in str(e)
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_get_consensus_label():
+    labels = data["labels"]
+
+    # getting consensus labels without pred_probs
+    consensus_label = get_majority_vote_label(labels)
+
+    # making synthetic data to test tiebreaks of get_consensus_label
+    # also testing passing labels as np.ndarray
+    labels_tiebreaks = np.array([[1, 2, 0], [1, 1, 0], [1, 0, 0], [2, 2, 2], [1, 2, 0], [1, 2, 0]])
+    pred_probs_tiebreaks = np.array(
+        [
+            [0.4, 0.4, 0.2],
+            [0.3, 0.6, 0.1],
+            [0.75, 0.2, 0.05],
+            [0.1, 0.4, 0.5],
+            [0.2, 0.4, 0.4],
+            [0.2, 0.4, 0.4],
+        ]
+    )
+    consensus_label = get_majority_vote_label(labels_tiebreaks, pred_probs_tiebreaks)
+
+    # more tiebreak testing (without pred_probs + non-overlapping annotators)
+    labels_tiebreaks = np.array(
+        [
+            [1, np.nan, np.nan, 2, np.nan],
+            [np.nan, 1, 0, np.nan, np.nan],
+            [np.nan, np.nan, 0, np.nan, np.nan],
+            [np.nan, 2, np.nan, np.nan, np.nan],
+            [2, np.nan, 0, 2, np.nan],
+            [np.nan, np.nan, np.nan, 2, 1],
+        ]
+    )
+    consensus_label = get_majority_vote_label(labels_tiebreaks)
+    assert all(consensus_label == np.array([1, 1, 0, 2, 2, 1]))
+
+
+def test_impute_nonoverlaping_annotators():
+    labels = np.array(
+        [
+            [1, np.nan, np.nan],
+            [np.nan, 1, 0],
+            [np.nan, 0, 0],
+            [np.nan, 2, 2],
+            [np.nan, 2, 0],
+            [np.nan, 2, 0],
+        ]
+    )
+    pred_probs = np.array(
+        [
+            [0.4, 0.4, 0.2],
+            [0.3, 0.6, 0.1],
+            [0.75, 0.2, 0.05],
+            [0.1, 0.4, 0.5],
+            [0.2, 0.4, 0.4],
+            [0.2, 0.4, 0.4],
+        ]
+    )
+
+    multiannotator_dict = get_label_quality_multiannotator(labels, pred_probs)
+    multiannotator_dict = get_label_quality_multiannotator(
+        labels, pred_probs, quality_method="agreement"
+    )
+
+
+def test_format_multiannotator_labels():
+    str_labels = np.array(
+        [
+            ["a", "b", "c"],
+            ["b", "b", np.nan],
+            ["z", np.nan, "c"],
+        ]
+    )
+    labels, label_map = format_multiannotator_labels(str_labels)
+
+    assert isinstance(labels, pd.DataFrame)
+    assert label_map[0] == "a"
+    assert label_map[3] == "z"
+
+    num_labels = pd.DataFrame(
+        [
+            [3, 2, 1],
+            [1, 2, np.nan],
+            [3, np.nan, 3],
+        ]
+    )
+    labels, label_map = format_multiannotator_labels(num_labels)
+
+
+def test_assert_valid_inputs_multiannotator_warnings(recwarn):
+    not_agree_labels = np.array([[2, 1, np.nan, 0], [1, 0, 2, np.nan]])
+    with pytest.warns(UserWarning, match="do not agree"):
+        assert_valid_inputs_multiannotator(not_agree_labels)
+
+    agree_labels = np.array([[1, 3, 3], [1, 4, 2]])
+    assert_valid_inputs_multiannotator(agree_labels)
+    # Assert no new warning were raised
+    assert len(recwarn) == 0

tests/test_multilabel_classification.py

@@ -0,0 +1,782 @@
+import itertools
+import typing
+
+import numpy as np
+import pytest
+import sklearn
+from hypothesis import given, settings
+from hypothesis import strategies as st
+from hypothesis.extra.numpy import arrays
+from hypothesis.strategies import composite
+from sklearn.linear_model import LogisticRegression
+from sklearn.multiclass import OneVsRestClassifier
+
+from cleanlab import multilabel_classification as ml_classification
+from cleanlab.internal import multilabel_scorer as ml_scorer
+from cleanlab.internal.multilabel_utils import get_onehot_num_classes, onehot2int, stack_complement
+from cleanlab.multilabel_classification import filter
+from cleanlab.multilabel_classification.dataset import (
+    common_multilabel_issues,
+    multilabel_health_summary,
+    overall_multilabel_health_score,
+    rank_classes_by_multilabel_quality,
+)
+from cleanlab.multilabel_classification.rank import get_label_quality_scores_per_class
+
+
+@pytest.fixture
+def labels():
+    return np.array(
+        [
+            [1, 0, 0],
+            [0, 1, 0],
+            [0, 0, 1],
+            [1, 1, 0],
+            [1, 0, 1],
+            [0, 1, 1],
+            [1, 1, 1],
+            [0, 0, 0],
+            [1, 0, 1],
+            [0, 1, 0],
+        ]
+    )
+
+
+@pytest.fixture
+def pred_probs_gold(labels):
+    pred_probs = np.array(
+        [
+            [0.203, 0.465, 0.612],
+            [0.802, 0.596, 0.43],
+            [0.776, 0.649, 0.391],
+            [0.201, 0.439, 0.633],
+            [0.203, 0.443, 0.584],
+            [0.814, 0.572, 0.332],
+            [0.201, 0.388, 0.544],
+            [0.778, 0.646, 0.392],
+            [0.796, 0.611, 0.387],
+            [0.199, 0.381, 0.58],
+        ]
+    )
+    assert pred_probs.shape == labels.shape
+    return pred_probs
+
+
+@pytest.fixture
+def pred_probs():
+    return np.array(
+        [
+            [0.9, 0.1, 0.2],
+            [0.5, 0.6, 0.4],
+            [0.75, 0.80, 0.85],
+            [0.9, 0.85, 0.2],
+            [0.9, 0.1, 0.85],
+            [0.5, 0.6, 0.85],
+            [0.9, 0.85, 0.85],
+            [0.8, 0.4, 0.2],
+            [0.9, 0.1, 0.85],
+            [0.15, 0.95, 0.05],
+        ]
+    )
+
+
+@pytest.fixture
+def pred_probs_multilabel():
+    return np.array(
+        [
+            [0.9, 0.1, 0.0, 0.4, 0.1],
+            [0.7, 0.8, 0.2, 0.3, 0.1],
+            [0.9, 0.8, 0.4, 0.2, 0.1],
+            [0.1, 0.1, 0.8, 0.3, 0.1],
+            [0.4, 0.5, 0.1, 0.1, 0.1],
+            [0.1, 0.1, 0.2, 0.1, 0.1],
+            [0.8, 0.1, 0.2, 0.1, 0.1],
+        ]
+    )
+
+
+@pytest.fixture
+def labels_multilabel():
+    return [[0], [0, 1], [0, 1], [2], [0, 2, 3], [], []]
+
+
+@pytest.fixture
+def data_multilabel(num_classes=5):
+    labels = []
+    pred_probs = []
+    for i in range(0, 100):
+        q = [0.1] * num_classes
+        pos = i % num_classes
+        labels.append([pos])
+        if i > 90:
+            pos = (pos + 2) % num_classes
+        q[pos] = 0.9
+        pred_probs.append(q)
+    return labels, np.array(pred_probs)
+
+
+@pytest.fixture
+def cv():
+    return sklearn.model_selection.StratifiedKFold(
+        n_splits=2,
+        shuffle=True,
+        random_state=42,
+    )
+
+
+@pytest.fixture
+def dummy_features(labels):
+    np.random.seed(42)
+    return np.random.rand(labels.shape[0], 2)
+
+
+def test_public_label_quality_scores(labels, pred_probs):
+    formatted_labels = onehot2int(labels)
+    assert isinstance(formatted_labels, list)
+    scores1 = ml_classification.get_label_quality_scores(formatted_labels, pred_probs)
+    assert len(scores1) == len(labels)
+    assert (scores1 >= 0).all() and (scores1 <= 1).all()
+    scores2 = ml_classification.get_label_quality_scores(
+        formatted_labels, pred_probs, method="confidence_weighted_entropy"
+    )
+    assert not np.isclose(scores1, scores2).all()
+    scores3 = ml_classification.get_label_quality_scores(
+        formatted_labels, pred_probs, adjust_pred_probs=True
+    )
+    assert not np.isclose(scores1, scores3).all()
+    scores4 = ml_classification.get_label_quality_scores(
+        formatted_labels,
+        pred_probs,
+        method="normalized_margin",
+        adjust_pred_probs=True,
+        aggregator_kwargs={"method": "exponential_moving_average"},
+    )
+    assert not np.isclose(scores1, scores4).all()
+    scores5 = ml_classification.get_label_quality_scores(
+        formatted_labels,
+        pred_probs,
+        method="normalized_margin",
+        adjust_pred_probs=True,
+        aggregator_kwargs={"method": "softmin"},
+    )
+    assert not np.isclose(scores4, scores5).all()
+    scores6 = ml_classification.get_label_quality_scores(
+        formatted_labels,
+        pred_probs,
+        method="normalized_margin",
+        adjust_pred_probs=True,
+        aggregator_kwargs={"method": "softmin", "temperature": 0.002},
+    )
+    assert not np.isclose(scores5, scores6).all()
+    scores7 = ml_classification.get_label_quality_scores(
+        formatted_labels,
+        pred_probs,
+        method="normalized_margin",
+        adjust_pred_probs=True,
+        aggregator_kwargs={"method": np.min},
+    )
+    assert np.isclose(scores6, scores7, rtol=1e-3).all()
+
+    with pytest.raises(ValueError) as e:
+        _ = ml_classification.get_label_quality_scores(
+            formatted_labels, pred_probs, method="badchoice"
+        )
+        assert "Invalid method name: badchoice" in str(e.value)
+
+    with pytest.raises(ValueError) as e:
+        _ = ml_classification.get_label_quality_scores(
+            formatted_labels, pred_probs, aggregator_kwargs={"method": "invalid"}
+        )
+        assert "Invalid aggregation method specified: 'invalid'" in str(e.value)
+
+
+class TestAggregator:
+    """Test the Aggregator class."""
+
+    @pytest.fixture
+    def base_scores(self):
+        return np.array([[0.6, 0.3, 0.7, 0.1, 0.9]])
+
+    @pytest.mark.parametrize(
+        "method",
+        [np.min, np.max, np.mean, np.median, "exponential_moving_average", "softmin"],
+        ids=lambda x: x.__name__ if callable(x) else str(x),
+    )
+    def test_aggregator_callable(self, method):
+        aggregator = ml_scorer.Aggregator(method=method)
+        assert callable(aggregator.method), "Aggregator should store a callable method"
+        assert callable(aggregator), "Aggregator should be callable"
+
+    @pytest.mark.parametrize(
+        "method,expected_score",
+        [
+            (np.min, 0.1),
+            (np.max, 0.9),
+            (np.mean, 0.52),
+            (np.median, 0.6),
+            ("exponential_moving_average", 0.436),
+            ("softmin", 0.128),
+        ],
+        ids=["min", "max", "mean", "median", "exponential_moving_average", "softmin"],
+    )
+    def test_aggregator_score(self, base_scores, method, expected_score):
+        aggregator = ml_scorer.Aggregator(method=method)
+        scores = aggregator(base_scores)
+        assert np.isclose(scores, np.array([expected_score]), rtol=1e-3).all()
+        assert scores.shape == (1,)
+
+    def test_invalid_method(self):
+        with pytest.raises(ValueError) as e:
+            _ = ml_scorer.Aggregator(method="invalid_method")
+            assert "Invalid aggregation method specified: 'invalid_method'" in str(
+                e.value
+            ), "String constructor has limited options"
+
+        with pytest.raises(TypeError) as e:
+            _ = ml_scorer.Aggregator(method=1)
+            assert "Expected callable method" in str(e.value), "Non-callable methods are not valid"
+
+    def test_invalid_score(self, base_scores):
+        aggregator = ml_scorer.Aggregator(method=np.min)
+        with pytest.raises(ValueError) as e:
+            _ = aggregator(base_scores[0])
+            assert "Expected 2D array" in str(e.value), "Aggregator expects 2D array"
+
+
+class TestMultilabelScorer:
+    """Test the MultilabelScorer class."""
+
+    @pytest.fixture
+    def docs_labels(self):
+        return np.array([[0, 1, 0], [1, 0, 1]])
+
+    @pytest.fixture
+    def docs_pred_probs(self):
+        return np.array([[0.1, 0.9, 0.7], [0.4, 0.1, 0.6]])
+
+    @pytest.fixture
+    def default_scorer(self):
+        return ml_scorer.MultilabelScorer()
+
+    @pytest.mark.parametrize(
+        "base_scorer", [scorer for scorer in ml_scorer.ClassLabelScorer], ids=lambda x: x.name
+    )
+    @pytest.mark.parametrize(
+        "aggregator", [np.min, np.max, np.mean, "exponential_moving_average", "softmin"]
+    )
+    @pytest.mark.parametrize("strict", [True, False], ids=["strict", ""])
+    def test_call(self, base_scorer, aggregator, strict, labels, pred_probs):
+        scorer = ml_scorer.MultilabelScorer(base_scorer, aggregator, strict=strict)
+        assert callable(scorer)
+
+        test_scores = scorer(labels, pred_probs)
+        assert isinstance(test_scores, np.ndarray)
+        assert test_scores.shape == (labels.shape[0],)
+
+        # Test base_scorer_kwargs
+        base_scorer_kwargs = {"adjust_pred_probs": True}
+        if scorer.base_scorer is not ml_scorer.ClassLabelScorer.CONFIDENCE_WEIGHTED_ENTROPY:
+            test_scores = scorer(labels, pred_probs, base_scorer_kwargs=base_scorer_kwargs)
+            assert isinstance(test_scores, np.ndarray)
+            assert test_scores.shape == (labels.shape[0],)
+        else:
+            with pytest.raises(ValueError) as e:
+                scorer(labels, pred_probs, base_scorer_kwargs=base_scorer_kwargs)
+                assert "adjust_pred_probs is not currently supported for" in str(e)
+
+    @pytest.mark.parametrize(
+        "base_scorer", [scorer for scorer in ml_scorer.ClassLabelScorer], ids=lambda x: x.name
+    )
+    def test_aggregate_kwargs(self, base_scorer):
+        """Make sure the instatiated aggregator kwargs can be overridden.
+        I.e. switching from a forgetting-factor 1.0 to 0.5.
+        """
+        class_label_quality_scores = np.array([[0.9, 0.9, 0.3], [0.4, 0.9, 0.6]])
+        aggregator = ml_scorer.Aggregator(ml_scorer.exponential_moving_average, alpha=1.0)
+        scorer = ml_scorer.MultilabelScorer(
+            base_scorer=base_scorer,
+            aggregator=aggregator,
+        )
+        scores = scorer.aggregate(class_label_quality_scores)
+        assert np.allclose(scores, np.array([0.3, 0.4]))
+        # Use different alpha, should change scores
+        new_scores = scorer.aggregate(class_label_quality_scores, alpha=0.0)
+        assert np.allclose(new_scores, np.array([0.9, 0.9]))
+
+    def test_get_class_label_quality_scores(self, default_scorer, docs_labels, docs_pred_probs):
+        """Test the get_class_label_quality_scores method."""
+        class_label_quality_scores = default_scorer.get_class_label_quality_scores(
+            docs_labels, docs_pred_probs
+        )
+        assert np.allclose(class_label_quality_scores, np.array([[0.9, 0.9, 0.3], [0.4, 0.9, 0.6]]))
+
+
+@pytest.mark.parametrize(
+    "method", ["self_confidence", "normalized_margin", "confidence_weighted_entropy"]
+)
+def test_class_label_scorer_from_str(method):
+    for m in (method, method.upper()):
+        scorer = ml_scorer.ClassLabelScorer.from_str(m)
+        assert callable(scorer)
+        with pytest.raises(ValueError):
+            ml_scorer.ClassLabelScorer.from_str(m.replace("_", "-"))
+
+
+@pytest.fixture
+def scorer():
+    return ml_scorer.MultilabelScorer(
+        base_scorer=ml_scorer.ClassLabelScorer.SELF_CONFIDENCE,
+        aggregator=np.min,
+    )
+
+
+def test_is_multilabel(labels):
+    assert ml_scorer._is_multilabel(labels)
+    assert not ml_scorer._is_multilabel(labels[:, 0])
+
+
+@pytest.mark.parametrize("class_names", [None, ["Apple", "Cat", "Dog", "Peach", "Bird"]])
+def test_common_multilabel_issues(class_names, pred_probs_multilabel, labels_multilabel):
+    df = common_multilabel_issues(
+        labels=labels_multilabel, pred_probs=pred_probs_multilabel, class_names=class_names
+    )
+    expected_issue_probabilities = [
+        0.14285714285714285,
+        0.14285714285714285,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+    ]
+    assert len(df) == 10
+    assert np.isclose(np.array(expected_issue_probabilities), df["Issue Probability"]).all()
+    if class_names:
+        expected_res = [
+            "Apple",
+            "Dog",
+            "Apple",
+            "Cat",
+            "Cat",
+            "Dog",
+            "Peach",
+            "Peach",
+            "Bird",
+            "Bird",
+        ]
+        assert list(df["Class Name"]) == expected_res
+    else:
+        assert "Class Name" not in df.columns
+
+
+def test_multilabel_find_label_issues(data_multilabel):
+    labels, pred_probs = data_multilabel
+    issues = filter.find_label_issues(
+        labels=labels, pred_probs=pred_probs, return_indices_ranked_by="self_confidence"
+    )
+    issues_lm = filter.find_label_issues(
+        labels, pred_probs, low_memory=True, return_indices_ranked_by="self_confidence"
+    )
+    intersection = len(list(set(issues).intersection(set(issues_lm))))
+    union = len(set(issues)) + len(set(issues_lm)) - intersection
+    assert float(intersection) / union > 0.95
+    # Check with return_indices_ranked_by=None
+    issues_mask = filter.find_label_issues(labels=labels, pred_probs=pred_probs)
+    issues_lm_mask = filter.find_label_issues(labels, pred_probs, low_memory=True)
+    issues_from_mask = np.where(issues_mask)[0]
+    issues_lm_from_mask = np.where(issues_lm_mask)[0]
+    intersection = len(list(set(issues_from_mask).intersection(set(issues_lm_from_mask))))
+    union = len(set(issues_from_mask)) + len(set(issues_lm_from_mask)) - intersection
+    assert float(intersection) / union > 0.95
+    # Check with low_memory=True, unused parameters rank_by_kwargs and n_jobs
+    rank_by_kwargs = {"adjust_pred_probs": None}
+    issues_lm2 = filter.find_label_issues(
+        labels,
+        pred_probs,
+        low_memory=True,
+        return_indices_ranked_by="self_confidence",
+        rank_by_kwargs=rank_by_kwargs,
+        n_jobs=1,
+    )
+    np.testing.assert_array_equal(issues_lm2, issues_lm)
+
+
+@pytest.mark.parametrize("min_examples_per_class", [10, 90])
+def test_multilabel_min_examples_per_class(data_multilabel, min_examples_per_class):
+    labels, pred_probs = data_multilabel
+    issues = filter.find_label_issues(
+        labels=labels, pred_probs=pred_probs, min_examples_per_class=min_examples_per_class
+    )
+    if min_examples_per_class == 10:
+        assert sum(issues) == 9
+    else:
+        assert sum(issues) == 0
+
+
+@pytest.mark.parametrize("num_to_remove_per_class", [None, [1, 1, 0, 0, 2], [1, 1, 0, 0, 1]])
+def test_multilabel_num_to_remove_per_class(data_multilabel, num_to_remove_per_class):
+    labels, pred_probs = data_multilabel
+
+    issues = filter.find_label_issues(
+        labels=labels, pred_probs=pred_probs, num_to_remove_per_class=num_to_remove_per_class
+    )
+    num_issues = sum(issues)
+    if num_to_remove_per_class is None:
+        assert num_issues == 9
+    else:
+        assert num_issues == sum(num_to_remove_per_class)
+
+
+@pytest.mark.parametrize("class_names", [None, ["Apple", "Cat", "Dog", "Peach", "Bird"]])
+def test_rank_classes_by_multilabel_quality(pred_probs_multilabel, labels_multilabel, class_names):
+    df_ranked = rank_classes_by_multilabel_quality(
+        pred_probs=pred_probs_multilabel, labels=labels_multilabel, class_names=class_names
+    )
+    expected_Label_Issues = [1, 0, 0, 0, 0]
+
+    expected_Label_Noise = [0.14285714285714285, 0.0, 0.0, 0.0, 0.0]
+
+    expected_Label_Quality_Score = [0.8571428571428572, 1.0, 1.0, 1.0, 1.0]
+
+    expected_Inverse_Label_Issues = [0, 1, 0, 0, 0]
+
+    expected_Inverse_Label_Noise = [0.0, 0.14285714285714285, 0.0, 0.0, 0.0]
+    assert list(df_ranked["Label Issues"]) == expected_Label_Issues
+
+    assert np.isclose(np.array(expected_Label_Noise), df_ranked["Label Noise"]).all()
+    assert np.isclose(
+        np.array(expected_Label_Quality_Score), df_ranked["Label Quality Score"]
+    ).all()
+    assert list(df_ranked["Inverse Label Issues"]) == expected_Inverse_Label_Issues
+    assert np.isclose(
+        np.array(expected_Inverse_Label_Noise), df_ranked["Inverse Label Noise"]
+    ).all()
+    if class_names:
+        expected_res = [
+            "Dog",
+            "Apple",
+            "Cat",
+            "Peach",
+            "Bird",
+        ]
+        assert list(df_ranked["Class Name"]) == expected_res
+    else:
+        assert "Class Name" not in df_ranked.columns
+
+
+def test_overall_multilabel_health_score(data_multilabel):
+    labels, pred_probs = data_multilabel
+    overall_label_health_score = overall_multilabel_health_score(
+        pred_probs=pred_probs, labels=labels
+    )
+    assert np.isclose(overall_label_health_score, 0.91)
+
+
+def test_get_class_label_quality_scores():
+    pred_probs = np.array(
+        [
+            [0.9, 0.1, 0.0, 0.4, 0.1],
+            [0.7, 0.8, 0.2, 0.3, 0.1],
+            [0.9, 0.8, 0.4, 0.2, 0.1],
+            [0.1, 0.1, 0.8, 0.3, 0.1],
+            [0.4, 0.5, 0.1, 0.1, 0.1],
+            [0.1, 0.1, 0.2, 0.1, 0.1],
+            [0.8, 0.1, 0.2, 0.1, 0.1],
+        ]
+    )
+    labels = [[0], [0, 1], [0, 1], [2], [0, 2, 3], [], []]
+    scores = get_label_quality_scores_per_class(pred_probs=pred_probs, labels=labels)
+    expected_res = [
+        [0.9, 0.9, 1.0, 0.6, 0.9],
+        [0.7, 0.8, 0.8, 0.7, 0.9],
+        [0.9, 0.8, 0.6, 0.8, 0.9],
+        [0.9, 0.9, 0.8, 0.7, 0.9],
+        [0.4, 0.5, 0.1, 0.1, 0.9],
+        [0.9, 0.9, 0.8, 0.9, 0.9],
+        [0.2, 0.9, 0.8, 0.9, 0.9],
+    ]
+    assert np.isclose(scores, np.array(expected_res)).all()
+
+
+def test_health_summary_multilabel(pred_probs_multilabel, labels_multilabel):
+    health_summary_multilabel = multilabel_health_summary(
+        pred_probs=pred_probs_multilabel, labels=labels_multilabel
+    )
+    expected_keys = [
+        "classes_by_multilabel_quality",
+        "common_multilabel_issues",
+        "overall_multilabel_health_score",
+    ]
+    assert sorted(health_summary_multilabel.keys()) == expected_keys
+
+
+@pytest.mark.parametrize(
+    "input",
+    [
+        [[0], [1, 2], [0, 2]],
+        [["a", "b"], ["b"]],
+        np.array([[[0, 1], [0, 1]], [[1, 1], [0, 0]]]),
+        1,
+    ],
+    ids=["lists of ids", "lists of strings", "3d array", "scalar"],
+)
+def test_is_multilabel_is_false(input):
+    assert not ml_scorer._is_multilabel(input)
+
+
+def test_stack_complement():
+    # Toy example
+    pred_probs_class = np.array([0.1, 0.9, 0.3, 0.8])
+    pred_probs_extended = stack_complement(pred_probs_class)
+    pred_probs_expected = np.array(
+        [
+            [0.9, 0.1],
+            [0.1, 0.9],
+            [0.7, 0.3],
+            [0.2, 0.8],
+        ]
+    )
+    assert np.isclose(pred_probs_extended, pred_probs_expected).all()
+
+    # Check preservation of probabilities
+    pred_probs_class = np.random.rand(100)
+    pred_probs_extended = stack_complement(pred_probs_class)
+    assert np.sum(pred_probs_extended, axis=1).all() == 1
+
+
+@pytest.mark.parametrize(
+    "pred_probs_test",
+    (None, "pred_probs"),
+    ids=["Without probabilities", "With probabilities"],
+)
+def test_get_onehot_num_classes(labels, pred_probs_test, request):
+    pred_probs_test = (
+        request.getfixturevalue(pred_probs_test)
+        if isinstance(pred_probs_test, str)
+        else pred_probs_test
+    )
+    labels_list = [np.nonzero(x)[0].tolist() for x in labels]
+    _, num_classes = get_onehot_num_classes(labels_list, pred_probs_test)
+    assert num_classes == 3
+
+
+def test_get_label_quality_scores_output(labels, pred_probs, scorer):
+    # Check that the function returns a dictionary with the correct keys.
+    scores = ml_scorer.get_label_quality_scores(labels, pred_probs, method=scorer)
+    assert isinstance(scores, np.ndarray)
+    assert scores.shape == (labels.shape[0],)
+    assert np.all(scores >= 0) and np.all(scores <= 1)
+    assert np.all(np.isfinite(scores))
+
+
+@pytest.mark.parametrize(
+    "given_labels,expected",
+    [
+        (
+            "labels",
+            np.full((3, 2), 0.5),
+        ),
+        (np.array([[0, 1], [0, 0], [1, 1]]), np.array([[2 / 3, 1 / 3], [1 / 3, 2 / 3]])),
+        (np.array([[0, 1], [0, 0], [0, 1], [0, 1]]), np.array([[4 / 4, 0 / 4], [1 / 4, 3 / 4]])),
+        (
+            np.array([[0, 1, 0, 0, 0, 0, 0, 0, 0]]),
+            np.array([[1, 0] if i != 1 else [0, 1] for i in range(9)]),
+        ),
+    ],
+    ids=[
+        "default",
+        "Missing class assignment configuration",
+        "Missing class",
+        "Handle more than 8 classes",
+    ],
+)
+def test_multilabel_py(given_labels, expected, request):
+    given_labels = (
+        request.getfixturevalue(given_labels) if isinstance(given_labels, str) else given_labels
+    )
+    py = ml_scorer.multilabel_py(given_labels)
+    assert isinstance(py, np.ndarray)
+    assert py.shape == (given_labels.shape[1], 2)
+    assert np.isclose(py, expected).all()
+
+
+@pytest.mark.parametrize("K", [2, 3, 4], ids=["K=2", "K=3", "K=4"])
+def test_get_split_generator(cv, K):
+    all_configurations = np.array(list(itertools.product([0, 1], repeat=K)))
+    given_labels = np.repeat(all_configurations, 2, axis=0)
+
+    split_generator = ml_scorer._get_split_generator(given_labels, cv)
+    assert isinstance(split_generator, typing.Generator)
+
+    train, test = next(split_generator)
+    for split in (train, test):
+        assert isinstance(split, np.ndarray)
+        assert np.isin(split, np.arange(given_labels.shape[0])).all()
+
+    # Test that the label distribution is relatively equal among the splits.
+    train_labels, test_labels = given_labels[train], given_labels[test]
+    _, train_counts = np.unique(train_labels, axis=0, return_counts=True)
+    _, test_counts = np.unique(test_labels, axis=0, return_counts=True)
+    # cv.get_n_splits() is 2, so we expect 1/2 of the labels in each split.
+    assert np.all(train_counts == 1)
+    assert np.all(test_counts == 1)
+
+
+# Test split_generator with rare/missing multilabel configurations
+@pytest.mark.parametrize("K", [2, 3, 4], ids=["K=2", "K=3", "K=4"])
+def test_get_split_generator_rare_configurations(cv, K):
+    all_configurations = np.array(list(itertools.product([0, 1], repeat=K)))
+    given_labels = np.repeat(all_configurations, 2, axis=0)
+
+    # Remove one configuration
+    given_labels = given_labels[~np.all(given_labels == all_configurations[0], axis=1)]
+
+    split_generator = ml_scorer._get_split_generator(given_labels, cv)
+    train, test = next(split_generator)
+    train_labels, test_labels = given_labels[train], given_labels[test]
+
+    # Test that the label distribution is relatively equal among the splits.
+    _, train_counts = np.unique(train_labels, axis=0, return_counts=True)
+    _, test_counts = np.unique(test_labels, axis=0, return_counts=True)
+    # cv.get_n_splits() is 2, so we expect 1/2 of the labels in each split.
+    assert np.all(train_counts == 1)
+    assert np.all(test_counts == 1)
+    assert len(train_counts) == len(test_counts) == len(all_configurations) - 1
+
+    # Remove one instance from labels
+    given_labels = given_labels[1:, :]
+
+    split_generator = ml_scorer._get_split_generator(given_labels, cv)
+    train, test = next(split_generator)
+    train_labels, test_labels = given_labels[train], given_labels[test]
+
+    # Test that the label distribution is relatively equal among the splits.
+    _, train_counts = np.unique(train_labels, axis=0, return_counts=True)
+    _, test_counts = np.unique(test_labels, axis=0, return_counts=True)
+    # cv.get_n_splits() is 2, so we expect 1/2 of the labels in each split,
+    # except for the class with one fewer instances.
+    assert len(train_counts) != len(test_counts)
+
+
+def test_get_cross_validated_multilabel_pred_probs(dummy_features, labels, cv, pred_probs_gold):
+    clf = OneVsRestClassifier(LogisticRegression(random_state=0))
+    pred_probs = ml_scorer.get_cross_validated_multilabel_pred_probs(
+        dummy_features,
+        labels,
+        clf=clf,
+        cv=cv,
+    )
+    assert isinstance(pred_probs, np.ndarray)
+    assert pred_probs.shape == labels.shape
+    assert np.all(pred_probs >= 0) and np.all(pred_probs <= 1)
+    assert np.all(np.isfinite(pred_probs))
+
+    # Gold master test - Ensure output is consistent
+    assert dummy_features.shape == (10, 2)
+    assert np.allclose(pred_probs, pred_probs_gold, atol=5e-4)
+
+
+class TestExponentialMovingAverage:
+    """Test the ml_scorer.expontential_moving_average function."""
+
+    @pytest.mark.parametrize("alpha", [0.5, None])
+    def test_valid_alpha(self, alpha):
+        # Test valid alpha values
+        for x, expected_ema in zip(
+            [
+                np.ones(5).reshape(1, -1),
+                np.array([[0.1, 0.2, 0.3]]),
+                np.array([x / 10 for x in range(1, 7)]).reshape(2, 3),
+            ],
+            [1, 0.175, np.array([0.175, 0.475])],
+        ):
+            ema = ml_scorer.exponential_moving_average(x, alpha=alpha)
+            assert np.allclose(ema, expected_ema, atol=1e-4)
+
+    @pytest.mark.parametrize(
+        "alpha,expected_ema",
+        [[0, 0.3], [1, 0.1]],
+        ids=["alpha=0", "alpha=1"],
+    )
+    def test_alpha_boundary(self, alpha, expected_ema):
+        # alpha = 0(1) should return the largest(smallest) value
+        X = np.array([[0.1, 0.2, 0.3]])
+        ema = ml_scorer.exponential_moving_average(X, alpha=alpha)
+        assert np.allclose(ema, expected_ema, atol=1e-4)
+
+    def test_invalid_alpha(self):
+        # Test that the exponential moving average raises an error
+        # when alpha is not in the interval [0, 1].
+        partial_error_msg = r"alpha must be in the interval \[0, 1\]"
+        for alpha in [-0.5, 1.5]:
+            with pytest.raises(ValueError, match=partial_error_msg):
+                ml_scorer.exponential_moving_average(np.ones(5).reshape(1, -1), alpha=alpha)
+
+
+def flip_labels(label, flip_prob):
+    """Flips binary labels with a given probability."""
+    rand_flip = np.random.choice(
+        [0, 1], size=label.shape, replace=True, p=[1 - flip_prob, flip_prob]
+    )
+    return np.abs(label - rand_flip).astype(int)
+
+
+@composite
+def cleanlab_data_strategy(draw):
+    num_classes = draw(st.integers(min_value=2, max_value=3))
+    num_samples = draw(st.integers(min_value=10, max_value=50))
+
+    # Generate true labels as one-hot encoded vectors for multi-label
+    true_labels = draw(
+        arrays(dtype=np.int8, shape=(num_samples, num_classes), elements=st.integers(0, 1))
+    )
+
+    # Generate noise matrix for multi-label and flip those values
+    flip_prob = 0.2
+    noisy_labels = flip_labels(true_labels, flip_prob)
+
+    # Multilabel find_issues raises a ValueError if all values are the same
+    # To avoid that we flip the first two values if all values are equal.
+    for i in range(noisy_labels.shape[1]):
+        if np.all(noisy_labels[:, i] == noisy_labels[0, i]):
+            noisy_labels[:2, i] = 1 - noisy_labels[:2, i]
+
+    # Generate predicted probabilities for each class for each sample
+    pred_probs = draw(
+        arrays(
+            dtype=np.float32,
+            shape=(num_samples, num_classes),
+            elements=st.floats(min_value=0, max_value=1, width=32),  # Specify width here
+        )
+    )
+
+    for i in range(num_samples):
+        for j in range(num_classes):
+            if draw(st.floats(min_value=0, max_value=1)) < 0.1:
+                # Set some probability values to exactly 0.5
+                pred_probs[i][j] = 0.5
+    return true_labels, noisy_labels, np.array(pred_probs)
+
+
+class TestMultiLabel:
+    @given(cleanlab_data_strategy())
+    @settings(deadline=20000)
+    def test_find_label_issues(self, data):
+        true_labels, noisy_labels, pred_probs = data
+        noisy_labels_list = onehot2int(noisy_labels)
+        is_issue = filter.find_label_issues(
+            labels=noisy_labels_list, pred_probs=np.array(pred_probs), n_jobs=1
+        )
+        threshold = 0.5
+        predicted_labels = (pred_probs >= threshold).astype(int)
+
+        # Check if predicted labels are the same as noisy labels for each example
+        labels_match = np.all(predicted_labels == noisy_labels, axis=1)
+
+        # For any example flagged as having an issue, there should be at least one label mismatch
+        assert not np.any(
+            is_issue & labels_match
+        ), "Examples with issues must have at least one label mismatch."

tests/test_noise_generation.py

@@ -0,0 +1,257 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+import numpy as np
+from cleanlab.benchmarking import noise_generation
+import pytest
+
+seed = 0
+np.random.seed(0)
+
+
+def test_main_pipeline(
+    verbose=False,
+    n=10,
+    valid_noise_matrix=True,
+    frac_zero_noise_rates=0,
+):
+    trace = 1.5
+    py = [0.1, 0.1, 0.2, 0.6]
+    K = len(py)
+    y = [z for i, p in enumerate(py) for z in [i] * int(p * n)]
+    nm = noise_generation.generate_noise_matrix_from_trace(
+        K=K,
+        trace=trace,
+        py=py,
+        seed=0,
+        valid_noise_matrix=valid_noise_matrix,
+        frac_zero_noise_rates=frac_zero_noise_rates,
+    )
+    # Check that trace is what its supposed to be
+    assert abs(trace - np.trace(nm) < 1e-2)
+    # Check that sum of probabilities is K
+    assert abs(nm.sum() - K) < 1e-4
+    # Check that sum of each column is 1
+    assert all(abs(nm.sum(axis=0) - 1) < 1e-4)
+    # Check that joint sums to 1.
+    assert abs(np.sum(nm * py) - 1 < 1e-4)
+    s = noise_generation.generate_noisy_labels(y, nm)
+    assert noise_generation.noise_matrix_is_valid(nm, py)
+
+
+def test_main_pipeline_fraczero_high():
+    test_main_pipeline(n=1000, frac_zero_noise_rates=0.75)
+
+
+def test_main_pipeline_verbose(verbose=True, n=10):
+    test_main_pipeline(verbose=verbose, n=n)
+
+
+def test_main_pipeline_many(verbose=False, n=1000):
+    test_main_pipeline(verbose=verbose, n=n)
+
+
+def test_main_pipeline_many_verbose_valid(verbose=True, n=100):
+    test_main_pipeline(verbose, n, valid_noise_matrix=True)
+
+
+def test_main_pipeline_many_valid(verbose=False, n=100):
+    test_main_pipeline(verbose, n, valid_noise_matrix=True)
+
+
+def test_main_pipeline_many_verbose(verbose=True, n=1000):
+    test_main_pipeline(verbose=verbose, n=n)
+
+
+@pytest.mark.parametrize("verbose", [True, False])
+def test_invalid_inputs_verify(verbose):
+    nm = np.array(
+        [
+            [0.2, 0.5],
+            [0.8, 0.5],
+        ]
+    )
+    py = [0.1, 0.8]
+    assert not noise_generation.noise_matrix_is_valid(nm, py, verbose=verbose)
+
+    nm = np.array(
+        [
+            [0.2, 0.5],
+            [0.8, 0.4],
+        ]
+    )
+    py = [0.1, 0.9]
+    assert not noise_generation.noise_matrix_is_valid(nm, py)
+
+    py = [0.1, 0.8]
+    assert not noise_generation.noise_matrix_is_valid(nm, py)
+
+
+def test_invalid_matrix():
+    nm = np.array(
+        [
+            [0.1, 0.9],
+            [0.9, 0.1],
+        ]
+    )
+    py = [0.1, 0.9]
+    assert not noise_generation.noise_matrix_is_valid(nm, py)
+
+
+def test_trace_less_than_1_error(trace=0.5):
+    try:
+        noise_generation.generate_noise_matrix_from_trace(3, trace)
+    except ValueError as e:
+        assert "trace > 1" in str(e)
+        with pytest.raises(ValueError) as e:
+            noise_generation.generate_noise_matrix_from_trace(3, trace)
+
+
+def test_trace_equals_1_error(trace=1):
+    test_trace_less_than_1_error(trace)
+
+
+def test_valid_no_py_error():
+    try:
+        noise_generation.generate_noise_matrix_from_trace(
+            K=3,
+            trace=2,
+            valid_noise_matrix=True,
+        )
+    except ValueError as e:
+        assert "py must be" in str(e)
+        with pytest.raises(ValueError) as e:
+            noise_generation.generate_noise_matrix_from_trace(
+                K=3,
+                trace=2,
+                valid_noise_matrix=True,
+            )
+
+
+def test_one_class_error():
+    try:
+        noise_generation.generate_noise_matrix_from_trace(
+            K=1,
+            trace=2,
+        )
+    except ValueError as e:
+        assert "must be >= 2" in str(e)
+        with pytest.raises(ValueError) as e:
+            noise_generation.generate_noise_matrix_from_trace(
+                K=1,
+                trace=1,
+            )
+
+
+def test_two_class_nofraczero():
+    trace = 1.1
+    nm = noise_generation.generate_noise_matrix_from_trace(
+        K=2,
+        trace=trace,
+        valid_noise_matrix=True,
+    )
+    assert not np.any(nm == 0)  # Make sure there is not a zero noise rate.
+    assert abs(trace - np.trace(nm) < 1e-2)
+
+
+def test_two_class_fraczero_high(valid=False):
+    trace = 1.8
+    frac_zero_noise_rates = 0.75
+    nm = noise_generation.generate_noise_matrix_from_trace(
+        K=2,
+        trace=trace,
+        valid_noise_matrix=valid,
+        frac_zero_noise_rates=frac_zero_noise_rates,
+    )
+    assert np.any(nm == 0)  # Make sure there is a zero noise rate.
+    assert abs(trace - np.trace(nm) < 1e-2)
+
+
+def test_two_class_fraczero_high_valid():
+    test_two_class_fraczero_high(True)
+
+
+def test_gen_probs_sum_empty():
+    f = noise_generation.generate_n_rand_probabilities_that_sum_to_m
+    assert len(f(n=0, m=1)) == 0
+
+
+def test_gen_probs_max_error():
+    f = noise_generation.generate_n_rand_probabilities_that_sum_to_m
+    try:
+        f(n=5, m=1, max_prob=0.1)
+    except ValueError as e:
+        assert "max_prob must be greater" in str(e)
+        with pytest.raises(ValueError) as e:
+            f(n=5, m=1, max_prob=0.1)
+
+
+def test_gen_probs_min_error():
+    f = noise_generation.generate_n_rand_probabilities_that_sum_to_m
+    try:
+        f(n=5, m=1, min_prob=0.9)
+    except ValueError as e:
+        assert "min_prob must be less" in str(e)
+        with pytest.raises(ValueError) as e:
+            f(n=5, m=1, min_prob=0.9)
+
+
+def test_probs_min_max_error():
+    f = noise_generation.generate_n_rand_probabilities_that_sum_to_m
+    min_prob = 0.5
+    max_prob = 0.5
+    try:
+        f(n=2, m=1, min_prob=min_prob, max_prob=max_prob)
+    except ValueError as e:
+        assert "min_prob must be less than max_prob" in str(e)
+        with pytest.raises(ValueError) as e:
+            f(n=5, m=1, min_prob=min_prob, max_prob=max_prob)
+
+
+def test_balls_zero():
+    f = noise_generation.randomly_distribute_N_balls_into_K_bins
+    K = 3
+    result = f(N=0, K=K)
+    assert len(result) == K
+    assert sum(result) == 0
+
+
+def test_balls_params():
+    f = noise_generation.randomly_distribute_N_balls_into_K_bins
+    N = 10
+    K = 10
+    for mx in [None, 1, 2, 3]:
+        for mn in [None, 1, 2, 3]:
+            r = f(
+                N=N,
+                K=K,
+                max_balls_per_bin=mx,
+                min_balls_per_bin=mn,
+            )
+            assert sum(r) == K
+            assert min(r) <= (K if mn is None else mn)
+            assert len(r) == K
+
+
+def test_max_iter():
+    trace = 2
+    K = 3
+    py = [1 / float(K)] * K
+    nm = noise_generation.generate_noise_matrix_from_trace(
+        K=K,
+        trace=trace,
+        valid_noise_matrix=True,
+        max_iter=1,
+        py=py,
+        seed=1,
+    )
+    assert abs(np.trace(nm) - trace) < 1e-6
+    assert abs(sum(np.dot(nm, py)) - 1) < 1e-6
+    nm2 = noise_generation.generate_noise_matrix_from_trace(
+        K=3,
+        trace=trace,
+        valid_noise_matrix=True,
+        py=[0.1, 0.1, 0.8],
+        max_iter=0,
+    )
+    assert nm2 is None

tests/test_outlier.py

@@ -0,0 +1,703 @@
+import numpy as np
+import pandas as pd
+import pytest
+from hypothesis import example, given, settings
+from hypothesis import strategies as st
+from sklearn.linear_model import LogisticRegression as LogReg
+from sklearn.neighbors import NearestNeighbors
+
+from cleanlab import count, outlier
+from cleanlab.benchmarking.noise_generation import (
+    generate_noise_matrix_from_trace,
+    generate_noisy_labels,
+)
+from cleanlab.count import get_confident_thresholds
+from cleanlab.internal.label_quality_utils import get_normalized_entropy
+from cleanlab.outlier import OutOfDistribution
+
+
+def make_data(
+    means=[[3, 2], [7, 7], [0, 8]],
+    covs=[[[5, -1.5], [-1.5, 1]], [[1, 0.5], [0.5, 4]], [[5, 1], [1, 5]]],
+    sizes=[80, 40, 40],
+    avg_trace=0.8,
+    seed=1,  # set to None for non-reproducible randomness
+):
+    np.random.seed(seed=seed)
+
+    m = len(means)  # number of classes
+    n = sum(sizes)
+    local_data = []
+    labels = []
+    test_data = []
+    test_labels = []
+
+    for idx in range(m):
+        local_data.append(
+            np.random.multivariate_normal(mean=means[idx], cov=covs[idx], size=sizes[idx])
+        )
+        test_data.append(
+            np.random.multivariate_normal(mean=means[idx], cov=covs[idx], size=sizes[idx])
+        )
+        labels.append(np.array([idx for i in range(sizes[idx])]))
+        test_labels.append(np.array([idx for i in range(sizes[idx])]))
+    X_train = np.vstack(local_data)
+    true_labels_train = np.hstack(labels)
+    X_test = np.vstack(test_data)
+    true_labels_test = np.hstack(test_labels)
+
+    # Compute p(true_label=k)
+    py = np.bincount(true_labels_train) / float(len(true_labels_train))
+
+    noise_matrix = generate_noise_matrix_from_trace(
+        m,
+        trace=avg_trace * m,
+        py=py,
+        valid_noise_matrix=True,
+        seed=seed,
+    )
+
+    # Generate our noisy labels using the noise_matrix.
+    s = generate_noisy_labels(true_labels_train, noise_matrix)
+    ps = np.bincount(s) / float(len(s))
+
+    # Compute inverse noise matrix
+    inv = count.compute_inv_noise_matrix(py, noise_matrix, ps=ps)
+
+    # Estimate pred_probs
+    latent = count.estimate_py_noise_matrices_and_cv_pred_proba(
+        X=X_train,
+        labels=s,
+        cv_n_folds=3,
+    )
+
+    label_errors_mask = s != true_labels_train
+
+    return {
+        "X_train": X_train,
+        "true_labels_train": true_labels_train,
+        "X_test": X_test,
+        "true_labels_test": true_labels_test,
+        "labels": s,
+        "label_errors_mask": label_errors_mask,
+        "ps": ps,
+        "py": py,
+        "noise_matrix": noise_matrix,
+        "inverse_noise_matrix": inv,
+        "est_py": latent[0],
+        "est_nm": latent[1],
+        "est_inv": latent[2],
+        "cj": latent[3],
+        "pred_probs": latent[4],
+        "m": m,
+        "n": n,
+    }
+
+
+# Global to be used by all test methods. Only compute this once for speed.
+data = make_data()
+
+
+def test_class_wrong_info_assert_valid_inputs():
+    features = data["X_train"]
+    pred_probs = data["pred_probs"]
+
+    OOD = OutOfDistribution()
+
+    # TESTING: _assert_valid_inputs() asserts correct errors in fit
+    try:
+        OOD.fit()
+    except Exception as e:
+        assert "Not enough information to compute scores" in str(e)
+        with pytest.raises(ValueError) as e:
+            OOD.fit()
+    try:
+        OOD.fit(features=features, pred_probs=pred_probs)
+    except Exception as e:
+        assert "Cannot fit to OOD Estimator to both features and pred_probs" in str(e)
+        with pytest.raises(ValueError) as e:
+            OOD.fit(features=features, pred_probs=pred_probs)
+
+    OOD = OutOfDistribution()
+
+    features_flat = np.ravel(features)
+    features_extra_dim = features[np.newaxis]
+    try:
+        OOD.fit(features=features_flat)
+    except Exception as e:
+        assert "array needs to be of shape (N, M)" in str(e)
+        with pytest.raises(ValueError) as e:
+            OOD.fit(features=features_flat)
+    try:
+        OOD.fit(features=features_extra_dim)
+    except Exception as e:
+        assert "array needs to be of shape (N, M)" in str(e)
+        with pytest.raises(ValueError) as e:
+            OOD.fit(features=features_extra_dim)
+
+    # TODO: DO WE NEED TO TESTING: _assert_valid_inputs() asserts correct errors in score?
+
+
+def test_class_wrong_info_fit_ood():
+    features = data["X_test"]
+    pred_probs = data["pred_probs"]
+    labels = data["labels"]
+
+    # TESTING: wrong param in params dict
+    try:
+        OOD = OutOfDistribution(params={"strange_param": -1})
+    except Exception as e:
+        assert "strange_param" in str(e)
+        with pytest.raises(ValueError) as e:
+            OOD = OutOfDistribution(params={"strange_param": -1})
+
+    #### SCORE wrong info
+
+    # TESTING: calling score before any fitting
+    OOD = OutOfDistribution()
+    try:
+        OOD.score(features=features)
+    except Exception as e:
+        assert "OOD estimator needs to be fit on features first" in str(e)
+        with pytest.raises(ValueError) as e:
+            OOD.score(features=features)
+
+    try:
+        OOD.score(pred_probs=pred_probs)
+    except Exception as e:
+        assert "OOD estimator needs to be fit on pred_probs first" in str(e)
+        with pytest.raises(ValueError) as e:
+            OOD.score(pred_probs=pred_probs)
+
+    # TESTING: calling scoring with opposite fitting
+    OOD_outlier = OutOfDistribution()
+    OOD_outlier.fit(features=features)
+    try:
+        OOD_outlier.score(pred_probs=pred_probs)
+    except Exception as e:
+        assert "OOD estimator needs to be fit on pred_probs first" in str(e)
+        with pytest.raises(ValueError) as e:
+            OOD_outlier.score(pred_probs=pred_probs)
+
+    OOD_ood = OutOfDistribution()
+    OOD_ood.fit(pred_probs=pred_probs, labels=labels)
+    try:
+        OOD_ood.score(features=features)
+    except Exception as e:
+        assert "OOD estimator needs to be fit on features first" in str(e)
+        with pytest.raises(ValueError) as e:
+            OOD_ood.score(features=features)
+
+
+def test_class_params_logic():
+    features = data["X_test"]
+    pred_probs = data["pred_probs"]
+
+    # TESTING: params dict is a copy
+    params_dict = {"k": 10, "t": 5}
+    OOD = OutOfDistribution(params=params_dict)
+
+    OOD.fit(features=features)
+    ood_params = OOD.params
+    params_dict = params_dict.update({"k": 20})
+    assert ood_params == OOD.params
+
+    # test calling functions with different params performs differently
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")  # Should be 7 warnings
+def test_class_public_func():
+    features = data["X_test"]
+    pred_probs = data["pred_probs"]
+    labels = data["true_labels_test"]
+
+    # Fit Logistic Regression model on X_train and estimate train pred_probs
+    logreg = LogReg(solver="lbfgs")
+    logreg.fit(data["X_train"], data["true_labels_train"])
+    train_pred_probs = logreg.predict_proba(data["X_train"])
+
+    # Get knn and confident_thresholds and pass them into OOD object initialization for testing already fitted logic
+    knn = NearestNeighbors(n_neighbors=7).fit(data["X_train"])
+    confident_thresholds = get_confident_thresholds(
+        pred_probs=train_pred_probs, labels=data["true_labels_train"]
+    )
+    OOD_outlier_already_fit = OutOfDistribution(params={"knn": knn})
+    OOD_ood_already_fit = OutOfDistribution(params={"confident_thresholds": confident_thresholds})
+
+    #### TESTING INITIALIZATION
+    # TESTING knn and confident_thresholds passed during initialization correctly
+    assert OOD_outlier_already_fit.params["knn"].n_neighbors == 7
+    assert (OOD_ood_already_fit.params["confident_thresholds"] == confident_thresholds).all()
+
+    #### TESTING FIT:
+    # Test fitting OOD object without labels and adjust_pred_probs=False
+    OOD_ood = OutOfDistribution(params={"adjust_pred_probs": False})
+    OOD_ood.fit(pred_probs=data["pred_probs"], labels=None)  # Warning
+    assert OOD_ood.params["adjust_pred_probs"] is False
+    OOD_ood.score(
+        pred_probs=pred_probs
+    )  # This should be ok without passing in labels since we are not adjusting
+
+    # Testing regular fit
+    OOD_ood = OutOfDistribution()
+    print(OOD_ood.params)
+    OOD_ood.fit(pred_probs=pred_probs, labels=labels)
+    print(OOD_ood.params)
+    OOD_outlier = OutOfDistribution()
+    OOD_outlier.fit(features=features)
+    print(OOD_outlier.params)
+    assert OOD_ood.params["confident_thresholds"] is not None and OOD_ood.params["knn"] is None
+    assert (
+        OOD_outlier.params["knn"] is not None and OOD_outlier.params["confident_thresholds"] is None
+    )
+    assert OOD_ood.params is not None and OOD_outlier.params is not None
+
+    # Testing calling fit on already fitted function (should not overwrite but warn)
+    OOD_outlier_already_fit.fit(features=features)  # Warning
+    assert OOD_outlier_already_fit.params["knn"].n_neighbors == 7  # Assert not overwritten
+    OOD_ood_already_fit.fit(pred_probs=pred_probs, labels=labels)  # Warning
+    assert (
+        OOD_ood_already_fit.params["confident_thresholds"] == confident_thresholds
+    ).all()  # Assert not overwritten
+
+    # Testing fit uses correct metrics given feature dimensionality
+    X_small = np.random.rand(20, 3)
+    OOD_euclidean = OutOfDistribution()
+    OOD_euclidean.fit(features=X_small)
+    # The metric attribute is the pairwise distance function implemented in scipy, use __name__ to get the name of the function
+    assert OOD_euclidean.params["knn"].metric.__name__ == "euclidean"
+    X_small_with_ood = np.vstack([X_small, [999999.0] * 3])
+    euclidean_score = OOD_euclidean.score(features=X_small_with_ood)
+    assert (np.max(euclidean_score) <= 1) and (np.min(euclidean_score) >= 0)
+    assert np.argmin(euclidean_score) == (euclidean_score.shape[0] - 1)
+
+    # Re-run tests with high dimensional dataset
+    X_large = np.hstack([np.zeros((200, 400)), np.random.rand(200, 1)])
+    OOD_cosine = OutOfDistribution()
+    OOD_cosine.fit(features=X_large)
+    assert OOD_cosine.params["knn"].metric == "cosine"
+    X_large_with_ood = np.vstack([X_large, [999999.0] * 401])
+    cosine_score = OOD_cosine.score(features=X_large_with_ood)
+    assert (np.max(cosine_score) <= 1) and (np.min(cosine_score) >= 0)
+    assert np.argmin(cosine_score) == (cosine_score.shape[0] - 1)
+
+    #### TESTING SCORE
+    ood_score = OOD_ood.score(pred_probs=pred_probs)
+    outlier_score = OOD_outlier.score(features=features)
+
+    assert ood_score is not None and outlier_score is not None
+    assert np.sum(ood_score) != np.sum(outlier_score)
+
+    #### TESTING FIT SCORE
+    OOD_ood_fs = OutOfDistribution()
+    ood_score_fs = OOD_ood_fs.fit_score(pred_probs=pred_probs, labels=labels)
+    OOD_outlier_fs = OutOfDistribution()
+    outlier_score_fs = OOD_outlier_fs.fit_score(features=features)
+
+    assert (
+        OOD_ood_fs.params["confident_thresholds"] is not None and OOD_ood_fs.params["knn"] is None
+    )
+    assert (
+        OOD_ood_fs.params["confident_thresholds"] == OOD_ood.params["confident_thresholds"]
+    ).all()
+
+    assert (
+        OOD_outlier.params["knn"] is not None and OOD_outlier.params["confident_thresholds"] is None
+    )
+    assert ood_score_fs is not None and outlier_score_fs is not None
+
+    assert np.sum(outlier_score_fs) - np.sum(outlier_score) < 1  # scores are similar
+    assert np.sum(ood_score_fs) - np.sum(ood_score) < 1  # scores are similar
+
+    # Testing calling fit_score on already fitted function
+    score_outlier_fs = OOD_outlier_already_fit.fit_score(features=features)  # Warning
+    assert OOD_outlier_already_fit.params["knn"].n_neighbors == 7
+    score_ood_fs = OOD_ood_already_fit.fit_score(pred_probs=pred_probs, labels=labels)  # Warning
+    assert (OOD_ood_already_fit.params["confident_thresholds"] == confident_thresholds).all()
+    assert (
+        score_outlier_fs is not None and score_ood_fs is not None
+    )  # Assert scores still calculated
+
+    # Testing calling fit_score repeatedly on already fitted function does not fit it more, just scores same as before
+    score_outlier_fs1 = OOD_outlier_already_fit.fit_score(features=features)  # Warning
+    score_ood_fs1 = OOD_ood_already_fit.fit_score(pred_probs=pred_probs, labels=labels)  # Warning
+    assert (score_outlier_fs == score_outlier_fs1).all() and (score_ood_fs1 == score_ood_fs).all()
+
+    # Testing scores calculated during already fitted fit_score identical to scores calculated during score.
+    score_outlier_s = OOD_outlier_already_fit.score(features=features)
+    score_ood_s = OOD_ood_already_fit.score(pred_probs=pred_probs)
+    assert (score_outlier_fs == score_outlier_s).all() and (score_ood_fs == score_ood_s).all()
+
+
+def test_get_ood_features_scores():
+    ood = OutOfDistribution()
+    X_train = data["X_train"]
+    X_test = data["X_test"]
+
+    # Create OOD datapoint
+    X_ood = np.array([[999999999.0, 999999999.0]])
+
+    # Add OOD datapoint to X_test
+    X_test_with_ood = np.vstack([X_test, X_ood])
+
+    # Fit nearest neighbors on X_train
+    knn = NearestNeighbors(n_neighbors=5, metric="euclidean").fit(X_train)
+    # Get KNN distance as outlier score
+    k = 5
+    knn_distance_to_score, _ = ood._get_ood_features_scores(features=X_test_with_ood, knn=knn, k=k)
+    # Checking that X_ood has the smallest outlier score among all the datapoints
+    assert np.argmin(knn_distance_to_score) == (knn_distance_to_score.shape[0] - 1)
+
+    # Get KNN distance as outlier score without passing k
+    # By default k=10 is used or k = n_neighbors when k > n_neighbors extracted from the knn
+    knn_distance_to_score, _ = ood._get_ood_features_scores(features=X_test_with_ood, knn=knn)
+    # Checking that X_ood has the smallest outlier score among all the datapoints
+    assert np.argmin(knn_distance_to_score) == (knn_distance_to_score.shape[0] - 1)
+
+    # Get KNN distance as outlier score passing k and t > 1
+    large_t_knn_distance_to_score, _ = ood._get_ood_features_scores(
+        features=X_test_with_ood, knn=knn, k=k, t=5
+    )
+
+    # Checking that X_ood has the smallest outlier score among all the datapoints
+    assert np.argmin(large_t_knn_distance_to_score) == (large_t_knn_distance_to_score.shape[0] - 1)
+
+    # Get KNN distance as outlier score passing k and t < 1
+    small_t_knn_distance_to_score, _ = ood._get_ood_features_scores(
+        features=X_test_with_ood, knn=knn, k=k, t=0.002
+    )
+
+    # Checking that X_ood has the smallest outlier score among all the datapoints
+    assert np.argmin(small_t_knn_distance_to_score) == (small_t_knn_distance_to_score.shape[0] - 1)
+    assert np.sum(small_t_knn_distance_to_score) >= np.sum(large_t_knn_distance_to_score)
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_default_k_and_model_get_ood_features_scores():
+    # Testing using 'None' as model param and correct setting of default k as max_k
+
+    # Create dataset with OOD example
+    X = data["X_test"]
+    X_ood = np.array([[999999999.0, 999999999.0]])
+    X_with_ood = np.vstack([X, X_ood])
+
+    instantiated_k = 10
+
+    # Create NN class object with small instantiated k and fit on data
+    knn = NearestNeighbors(n_neighbors=instantiated_k, metric="euclidean").fit(X_with_ood)
+
+    ood = OutOfDistribution()
+
+    avg_knn_distances_default_model, _ = ood._get_ood_features_scores(
+        features=X_with_ood,
+        k=instantiated_k,  # this should use default estimator (same as above) and k = instantiated_k
+    )
+
+    avg_knn_distances_default_k, knn2 = ood._get_ood_features_scores(
+        features=X_with_ood,  # default k should be set to 10 == instantiated_k
+    )
+    assert isinstance(knn2, type(knn))
+
+    avg_knn_distances, _ = ood._get_ood_features_scores(
+        features=None,
+        knn=knn,
+        k=25,  # this should throw user warn, k should be set to instantiated_k
+    )
+
+    # Score sums should be equal because the three estimators used have identical params and fit
+    assert avg_knn_distances.sum() == avg_knn_distances_default_model.sum()
+    assert avg_knn_distances_default_k.sum() == avg_knn_distances.sum()
+
+    avg_knn_distances_large_k, _ = ood._get_ood_features_scores(
+        features=X_with_ood,
+        k=25,  # this should use default estimator and k = 25
+    )
+
+    avg_knn_distances_tiny_k, _ = ood._get_ood_features_scores(
+        features=None,
+        knn=knn,
+        k=1,  # this should use knn estimator and k = 1
+    )
+
+    avg_knn_distances_tiny_k_default, _ = ood._get_ood_features_scores(
+        features=X_with_ood,
+        k=1,  # this should use default estimator and k = 1
+    )
+
+    # Score sums should be different because k = user param for estimators and k != 10.
+    assert avg_knn_distances_tiny_k.sum() != avg_knn_distances.sum()
+    assert avg_knn_distances_large_k.sum() != avg_knn_distances.sum()
+    assert avg_knn_distances_tiny_k_default.sum() != avg_knn_distances_default_model.sum()
+
+    # Test that when knn is None ValueError raised if passed in k > len(features)
+    try:
+        ood._get_ood_features_scores(
+            features=X_with_ood,
+            knn=None,
+            k=len(X_with_ood) + 1,  # this should throw ValueError, k ! > len(features)
+        )
+    except Exception as e:
+        assert "nearest neighbors" in str(e)
+        with pytest.raises(ValueError) as e:
+            ood._get_ood_features_scores(
+                features=X_with_ood,
+                knn=None,
+                k=len(X_with_ood) + 1,  # this should throw ValueError, k ! > len(features)
+            )
+
+
+def test_not_enough_info_get_ood_features_scores():
+    # Testing calling function with not enough information to calculate outlier scores
+    ood = OutOfDistribution()
+    try:
+        ood._get_ood_features_scores(
+            features=None,
+            knn=None,  # this should throw TypeError because knn=None and features=None
+        )
+    except Exception as e:
+        assert "Both knn and features arguments" in str(e)
+        with pytest.raises(ValueError) as e:
+            ood._get_ood_features_scores(
+                features=None,
+                knn=None,  # this should throw TypeError because knn=None and features=None
+            )
+
+
+def test_ood_predictions_scores():
+    # Create and add OOD datapoint to test set
+    X = data["X_test"]
+    means = [[3, 2], [7, 7], [0, 8]]
+    X_ood = np.array(means).mean(axis=0)
+    X_with_ood = np.vstack([X, X_ood])
+
+    y = data["true_labels_test"]
+    y_with_ood = np.hstack([y, data["true_labels_train"][1]])
+
+    # Fit Logistic Regression model on X_train and estimate pred_probs
+    logreg = LogReg(solver="lbfgs")
+    logreg.fit(data["X_train"], data["true_labels_train"])
+    pred_probs = logreg.predict_proba(X_with_ood)
+
+    ### Test non-adjusted OOD score logic
+    ood_predictions_scores_entropy, _ = outlier._get_ood_predictions_scores(
+        pred_probs=pred_probs,
+        adjust_pred_probs=False,
+    )
+
+    # adjust pred probs should be False by default
+    ood_predictions_scores_least_confidence, _ = outlier._get_ood_predictions_scores(
+        pred_probs=pred_probs,
+        method="least_confidence",
+        adjust_pred_probs=False,
+    )
+
+    ood_predictions_scores_gen, _ = outlier._get_ood_predictions_scores(
+        pred_probs=pred_probs,
+        method="gen",
+        adjust_pred_probs=False,
+        M=3,  # Totally three classes
+    )
+
+    # check OOD scores calculated correctly
+    assert (1.0 - get_normalized_entropy(pred_probs) == ood_predictions_scores_entropy).all()
+    assert (pred_probs.max(axis=1) == ood_predictions_scores_least_confidence).all()
+    assert ood_predictions_scores_gen.max() < 1
+    assert ood_predictions_scores_gen.min() > 0
+    assert np.where(np.sort(ood_predictions_scores_entropy) == ood_predictions_scores_entropy[-1])[
+        0
+    ] < 0.02 * len(ood_predictions_scores_entropy)
+    assert np.where(
+        np.sort(ood_predictions_scores_least_confidence)
+        == ood_predictions_scores_least_confidence[-1]
+    )[0] < 0.02 * len(ood_predictions_scores_least_confidence)
+    assert np.where(np.sort(ood_predictions_scores_gen) == ood_predictions_scores_gen[-1])[
+        0
+    ] < 0.02 * len(ood_predictions_scores_gen)
+
+    ### Test adjusted OOD score logic
+    (
+        ood_predictions_scores_adj_entropy,
+        confident_thresholds_adj_entropy,
+    ) = outlier._get_ood_predictions_scores(
+        pred_probs=pred_probs,
+        labels=y_with_ood,
+        adjust_pred_probs=True,
+        method="entropy",
+    )
+
+    (
+        ood_predictions_scores_adj_least_confidence,
+        confident_thresholds_adj_least_confidence,
+    ) = outlier._get_ood_predictions_scores(
+        pred_probs=pred_probs,
+        labels=y_with_ood,
+        adjust_pred_probs=True,
+        method="least_confidence",
+    )
+
+    # test confident thresholds calculated correctly
+    confident_thresholds = get_confident_thresholds(
+        labels=y_with_ood, pred_probs=pred_probs, multi_label=False
+    )
+
+    assert (confident_thresholds == confident_thresholds_adj_entropy).all()
+    assert (confident_thresholds_adj_least_confidence == confident_thresholds_adj_entropy).all()
+
+    # check adjusted OOD scores different from non adjust OOD scores
+    assert not (ood_predictions_scores_adj_entropy == ood_predictions_scores_entropy).all()
+    assert not (
+        ood_predictions_scores_adj_least_confidence == ood_predictions_scores_least_confidence
+    ).all()
+
+    ### Test pre-calculated confident thresholds logic
+    ood_predictions_scores_2, confident_thresholds_2 = outlier._get_ood_predictions_scores(
+        pred_probs=pred_probs,
+        confident_thresholds=confident_thresholds,
+        adjust_pred_probs=True,
+    )
+
+    assert (confident_thresholds_2 == confident_thresholds).all()
+    assert (ood_predictions_scores_2 == ood_predictions_scores_adj_entropy).all()
+
+    # test using labels list type works
+    y_with_ood_list = y_with_ood.tolist()
+    (
+        ood_predictions_scores_adj_entropy_list,
+        confident_thresholds_adj_entropy_list,
+    ) = outlier._get_ood_predictions_scores(
+        pred_probs=pred_probs,
+        labels=y_with_ood_list,
+        adjust_pred_probs=True,
+        method="entropy",
+    )
+
+    # test using labels series type works
+    y_with_ood_series = pd.Series(y_with_ood)
+    (
+        ood_predictions_scores_adj_entropy_series,
+        confident_thresholds_adj_entropy_series,
+    ) = outlier._get_ood_predictions_scores(
+        pred_probs=pred_probs,
+        labels=y_with_ood_series,
+        adjust_pred_probs=True,
+        method="entropy",
+    )
+
+    assert (confident_thresholds_adj_entropy_list == confident_thresholds_adj_entropy).all()
+    assert (confident_thresholds_adj_entropy_series == confident_thresholds_adj_entropy).all()
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_wrong_info_get_ood_predictions_scores():
+    # Test calling function with not enough information to calculate ood scores
+    try:
+        outlier._get_ood_predictions_scores(
+            pred_probs=data["pred_probs"],
+            labels=None,
+            adjust_pred_probs=True,  # this should throw ValueError because knn=None and features=None
+        )
+    except Exception as e:
+        assert "Cannot calculate adjust_pred_probs without labels" in str(e)
+        with pytest.raises(ValueError) as e:
+            outlier._get_ood_predictions_scores(
+                pred_probs=data["pred_probs"],
+                labels=None,
+                adjust_pred_probs=True,  # this should throw ValueError because knn=None and features=None
+            )
+
+    # Test calling function with not enough information to calculate ood scores
+    try:
+        outlier._get_ood_predictions_scores(
+            pred_probs=data["pred_probs"],
+            adjust_pred_probs=True,  # this should throw ValueError because knn=None and features=None
+        )
+    except Exception as e:
+        assert "Cannot calculate adjust_pred_probs without labels" in str(e)
+        with pytest.raises(ValueError) as e:
+            outlier._get_ood_predictions_scores(
+                pred_probs=data["pred_probs"],
+                adjust_pred_probs=True,  # this should throw ValueError because not enough data provided
+            )
+
+    # Test calling function with not a real method
+    try:
+        outlier._get_ood_predictions_scores(
+            pred_probs=data["pred_probs"],
+            labels=data["labels"],
+            adjust_pred_probs=True,
+            method="not_a_real_method",  # this should throw ValueError because method not real method
+        )
+    except Exception as e:
+        assert "not a valid OOD scoring" in str(e)
+        with pytest.raises(ValueError) as e:
+            outlier._get_ood_predictions_scores(
+                pred_probs=data["pred_probs"],
+                labels=data["labels"],
+                adjust_pred_probs=True,
+                method="not_a_real_method",  # this should throw ValueError because method not real method
+            )
+
+    # Test calling function with too much information to calculate ood scores
+    outlier._get_ood_predictions_scores(
+        pred_probs=data["pred_probs"],
+        labels=data["labels"],
+        adjust_pred_probs=False,  # this should user warning because provided info is not used
+    )
+
+
+@given(
+    fill_value=st.floats(
+        min_value=5 * float(np.finfo(np.float64).eps),
+        max_value=5,
+        exclude_min=False,
+        allow_subnormal=False,
+        allow_infinity=False,
+        allow_nan=False,
+    ),
+    K=st.integers(min_value=2, max_value=100),
+)
+@example(K=1, fill_value=0.0)
+@settings(deadline=None)
+def test_scores_for_identical_examples(fill_value, K):
+    N = 100
+
+    features = np.full((N, K), fill_value=fill_value)
+    ood = OutOfDistribution()
+    scores = ood.fit_score(features=features, verbose=False)
+
+    # Dataset with only
+    expected_score = np.full(N, 1.0)
+    np.testing.assert_array_equal(
+        scores,
+        expected_score,
+        err_msg=f"The calculated distances were {ood.params['knn'].kneighbors()}",
+    )
+
+
+@given(K=st.integers(min_value=2, max_value=100))
+@settings(max_examples=10000, deadline=None)
+def test_scores_for_identical_examples_across_rows(K):
+    N = 100
+    fill_value = np.random.random(K)
+    features = np.full((N, K), fill_value=fill_value)
+    ood = OutOfDistribution()
+    scores = ood.fit_score(features=features, verbose=False)
+
+    # Dataset with only
+    expected_score = np.full(N, 1.0)
+    np.testing.assert_array_equal(
+        scores,
+        expected_score,
+        err_msg=f"The calculated distances were {ood.params['knn'].kneighbors()}",
+    )
+
+    if K < 4:
+        # This little changes should not affect euclidean calculation
+        features += np.random.random(features.shape) * 1e-10
+        ood = OutOfDistribution()
+        scores = ood.fit_score(features=features, verbose=False)
+        np.testing.assert_array_equal(
+            scores,
+            expected_score,
+            err_msg=f"The calculated distances were {ood.params['knn'].kneighbors()}",
+        )

tests/test_rank.py

@@ -0,0 +1,404 @@
+import numpy as np
+import pytest
+from cleanlab import rank
+from cleanlab.internal.label_quality_utils import _subtract_confident_thresholds
+from cleanlab.benchmarking.noise_generation import generate_noise_matrix_from_trace
+from cleanlab.benchmarking.noise_generation import generate_noisy_labels
+from cleanlab import count
+from cleanlab import outlier
+from sklearn.neighbors import NearestNeighbors
+
+
+def make_data(
+    means=[[3, 2], [7, 7], [0, 8]],
+    covs=[[[5, -1.5], [-1.5, 1]], [[1, 0.5], [0.5, 4]], [[5, 1], [1, 5]]],
+    sizes=[80, 40, 40],
+    avg_trace=0.8,
+    seed=1,  # set to None for non-reproducible randomness
+):
+    np.random.seed(seed=seed)
+
+    m = len(means)  # number of classes
+    n = sum(sizes)
+    local_data = []
+    labels = []
+    test_data = []
+    test_labels = []
+
+    for idx in range(m):
+        local_data.append(
+            np.random.multivariate_normal(mean=means[idx], cov=covs[idx], size=sizes[idx])
+        )
+        test_data.append(
+            np.random.multivariate_normal(mean=means[idx], cov=covs[idx], size=sizes[idx])
+        )
+        labels.append(np.array([idx for i in range(sizes[idx])]))
+        test_labels.append(np.array([idx for i in range(sizes[idx])]))
+    X_train = np.vstack(local_data)
+    true_labels_train = np.hstack(labels)
+    X_test = np.vstack(test_data)
+    true_labels_test = np.hstack(test_labels)
+
+    # Compute p(true_label=k)
+    py = np.bincount(true_labels_train) / float(len(true_labels_train))
+
+    noise_matrix = generate_noise_matrix_from_trace(
+        m,
+        trace=avg_trace * m,
+        py=py,
+        valid_noise_matrix=True,
+        seed=seed,
+    )
+
+    # Generate our noisy labels using the noise_matrix.
+    s = generate_noisy_labels(true_labels_train, noise_matrix)
+    ps = np.bincount(s) / float(len(s))
+
+    # Compute inverse noise matrix
+    inv = count.compute_inv_noise_matrix(py, noise_matrix, ps=ps)
+
+    # Estimate pred_probs
+    latent = count.estimate_py_noise_matrices_and_cv_pred_proba(
+        X=X_train,
+        labels=s,
+        cv_n_folds=3,
+    )
+
+    label_errors_mask = s != true_labels_train
+
+    return {
+        "X_train": X_train,
+        "true_labels_train": true_labels_train,
+        "X_test": X_test,
+        "true_labels_test": true_labels_test,
+        "labels": s,
+        "label_errors_mask": label_errors_mask,
+        "ps": ps,
+        "py": py,
+        "noise_matrix": noise_matrix,
+        "inverse_noise_matrix": inv,
+        "est_py": latent[0],
+        "est_nm": latent[1],
+        "est_inv": latent[2],
+        "cj": latent[3],
+        "pred_probs": latent[4],
+        "m": m,
+        "n": n,
+    }
+
+
+# Global to be used by all test methods. Only compute this once for speed.
+data = make_data()
+
+
+def test_get_normalized_margin_for_each_label():
+    scores = rank.get_normalized_margin_for_each_label(data["labels"], data["pred_probs"])
+    label_errors = np.arange(len(data["labels"]))[data["label_errors_mask"]]
+    least_confident_label = np.argmin(scores)
+    most_confident_label = np.argmax(scores)
+    assert least_confident_label in label_errors
+    assert most_confident_label not in label_errors
+
+
+def test_get_self_confidence_for_each_label():
+    scores = rank.get_self_confidence_for_each_label(data["labels"], data["pred_probs"])
+    label_errors = np.arange(len(data["labels"]))[data["label_errors_mask"]]
+    least_confident_label = np.argmin(scores)
+    most_confident_label = np.argmax(scores)
+    assert least_confident_label in label_errors
+    assert most_confident_label not in label_errors
+
+
+def test_bad_rank_by_parameter_error():
+    with pytest.raises(ValueError) as e:
+        _ = rank.order_label_issues(
+            label_issues_mask=data["label_errors_mask"],
+            labels=data["labels"],
+            pred_probs=data["pred_probs"],
+            rank_by="not_a_real_method",
+        )
+
+
+@pytest.mark.parametrize(
+    "scoring_method_func",
+    [
+        ("self_confidence", rank.get_self_confidence_for_each_label),
+        ("normalized_margin", rank.get_normalized_margin_for_each_label),
+        ("confidence_weighted_entropy", rank.get_confidence_weighted_entropy_for_each_label),
+    ],
+)
+@pytest.mark.parametrize("adjust_pred_probs", [False, True])
+def test_order_label_issues_using_scoring_func_ranking(scoring_method_func, adjust_pred_probs):
+    # test all scoring methods with the scoring function
+
+    method, scoring_func = scoring_method_func
+
+    # check if method supports adjust_pred_probs
+    # do not run the test below if the method does not support adjust_pred_probs
+    # confidence_weighted_entropy scoring method does not support adjust_pred_probs
+    if not (adjust_pred_probs == True and method == "confidence_weighted_entropy"):
+        indices = np.arange(len(data["label_errors_mask"]))[
+            data["label_errors_mask"]
+        ]  # indices of label issues
+
+        label_issues_indices = rank.order_label_issues(
+            label_issues_mask=data["label_errors_mask"],
+            labels=data["labels"],
+            pred_probs=data["pred_probs"],
+            rank_by=method,
+            rank_by_kwargs={"adjust_pred_probs": adjust_pred_probs},
+        )
+
+        # test scoring function with scoring method passed as arg
+        scores = rank.get_label_quality_scores(
+            data["labels"],
+            data["pred_probs"],
+            method=method,
+            adjust_pred_probs=adjust_pred_probs,
+        )
+        scores = scores[data["label_errors_mask"]]
+        score_idx = sorted(list(zip(scores, indices)), key=lambda y: y[0])  # sort indices by score
+        label_issues_indices2 = [z[1] for z in score_idx]
+        assert all(
+            label_issues_indices == label_issues_indices2
+        ), f"Test failed with scoring method: {method}"
+
+        # test individual scoring function
+        # only test if adjust_pred_probs=False because the individual scoring functions do not adjust pred_probs
+        if not adjust_pred_probs:
+            scores = scoring_func(data["labels"], data["pred_probs"])
+            scores = scores[data["label_errors_mask"]]
+            score_idx = sorted(
+                list(zip(scores, indices)), key=lambda y: y[0]
+            )  # sort indices by score
+            label_issues_indices3 = [z[1] for z in score_idx]
+            assert all(
+                label_issues_indices == label_issues_indices3
+            ), f"Test failed with scoring method: {method}"
+
+
+def test__subtract_confident_thresholds():
+    labels = data["labels"]
+    pred_probs = data["pred_probs"]
+
+    # subtract confident class thresholds and renormalize
+    pred_probs_adj = _subtract_confident_thresholds(labels, pred_probs)
+
+    assert (pred_probs_adj > 0).all()  # all pred_prob are positive numbers
+    assert (
+        abs(1 - pred_probs_adj.sum(axis=1)) < 1e-6
+    ).all()  # all pred_prob sum to 1 with some small precision error
+
+
+@pytest.mark.parametrize(
+    "method",
+    [
+        "self_confidence",
+        "normalized_margin",
+        "confidence_weighted_entropy",
+    ],
+)
+@pytest.mark.parametrize("adjust_pred_probs", [False, True])
+@pytest.mark.parametrize("weight_ensemble_members_by", ["uniform", "accuracy", "log_loss_search"])
+def test_ensemble_scoring_func(method, adjust_pred_probs, weight_ensemble_members_by):
+    labels = data["labels"]
+    pred_probs = data["pred_probs"]
+
+    # check if method supports adjust_pred_probs
+    # do not run the test below if the method does not support adjust_pred_probs
+    # confidence_weighted_entropy scoring method does not support adjust_pred_probs
+    if not (adjust_pred_probs == True and method == "confidence_weighted_entropy"):
+        # baseline scenario where all the pred_probs are the same in the ensemble list
+        num_repeat = 3
+        pred_probs_list = list(np.repeat([pred_probs], num_repeat, axis=0))
+
+        # get label quality score with single pred_probs
+        label_quality_scores = rank.get_label_quality_scores(
+            labels, pred_probs, method=method, adjust_pred_probs=adjust_pred_probs
+        )
+
+        # get ensemble label quality score
+        label_quality_scores_ensemble = rank.get_label_quality_ensemble_scores(
+            labels,
+            pred_probs_list,
+            method=method,
+            adjust_pred_probs=adjust_pred_probs,
+            weight_ensemble_members_by=weight_ensemble_members_by,
+        )
+
+        # if all pred_probs in the list are the same, then ensemble score should be the same as the regular score
+        # account for small precision error due to averaging of scores
+        assert (
+            abs(label_quality_scores - label_quality_scores_ensemble) < 1e-6
+        ).all(), f"Test failed with scoring method: {method}"
+
+
+def test_bad_weight_ensemble_members_by_parameter_error():
+    with pytest.raises(ValueError) as e:
+        labels = data["labels"]
+        pred_probs = data["pred_probs"]
+
+        # baseline scenario where all the pred_probs are the same in the ensemble list
+        num_repeat = 3
+        pred_probs_list = list(np.repeat([pred_probs], num_repeat, axis=0))
+
+        _ = rank.get_label_quality_ensemble_scores(
+            labels,
+            pred_probs_list,
+            weight_ensemble_members_by="not_a_real_method",  # this should raise ValueError
+        )
+
+
+def test_custom_weights():
+    with pytest.raises(AssertionError) as e:
+        labels = data["labels"]
+        pred_probs = data["pred_probs"]
+
+        # baseline scenario where all the pred_probs are the same in the ensemble list
+        num_repeat = 3
+        pred_probs_list = list(np.repeat([pred_probs], num_repeat, axis=0))
+
+        # baseline scenario where custom_weights are uniform
+        custom_weights = np.ones(num_repeat) / 3
+
+        scores_custom_weights = rank.get_label_quality_ensemble_scores(
+            labels,
+            pred_probs_list,
+            weight_ensemble_members_by="custom",
+            custom_weights=custom_weights,  # this should raise AssertionError
+        )
+
+        scores_uniform_weights = rank.get_label_quality_ensemble_scores(
+            labels, pred_probs_list, weight_ensemble_members_by="uniform"
+        )
+
+        # if custom_weights are uniform, then it should be the same as using weight_ensemble_members_by="uniform"
+        assert (scores_custom_weights == scores_uniform_weights).all()
+
+
+def test_empty_custom_weights_error():
+    labels = data["labels"]
+    pred_probs = data["pred_probs"]
+
+    # baseline scenario where all the pred_probs are the same in the ensemble list
+    num_repeat = 3
+    pred_probs_list = list(np.repeat([pred_probs], num_repeat, axis=0))
+
+    with pytest.raises(AssertionError) as e:
+        _ = rank.get_label_quality_ensemble_scores(
+            labels,
+            pred_probs_list,
+            weight_ensemble_members_by="custom",
+            custom_weights=None,  # this should raise AssertionError because custom_weights is None
+        )
+
+
+def test_wrong_length_custom_weights_error():
+    labels = data["labels"]
+    pred_probs = data["pred_probs"]
+
+    # baseline scenario where all the pred_probs are the same in the ensemble list
+    num_repeat = 3
+    pred_probs_list = list(np.repeat([pred_probs], num_repeat, axis=0))
+
+    # baseline scenario where custom_weights are uniform
+    custom_weights = np.ones(num_repeat) / 3
+
+    with pytest.raises(AssertionError) as e:
+        _ = rank.get_label_quality_ensemble_scores(
+            labels,
+            pred_probs_list,
+            weight_ensemble_members_by="custom",
+            custom_weights=custom_weights[1:],
+            # this should raise AssertionError because length of custom_weights don't match len(pred_probs_list)
+        )
+
+
+def test_wrong_weight_ensemble_members_by_for_custom_weights_error():
+    labels = data["labels"]
+    pred_probs = data["pred_probs"]
+
+    # baseline scenario where all the pred_probs are the same in the ensemble list
+    num_repeat = 3
+    pred_probs_list = list(np.repeat([pred_probs], num_repeat, axis=0))
+
+    # baseline scenario where custom_weights are uniform
+    custom_weights = np.ones(num_repeat) / 3
+
+    with pytest.raises(ValueError) as e:
+        _ = rank.get_label_quality_ensemble_scores(
+            labels,
+            pred_probs_list,
+            weight_ensemble_members_by="accuracy",
+            # this should raise ValueError because custom_weights array is provided
+            custom_weights=custom_weights,
+        )
+
+
+def test_bad_pred_probs_list_parameter_error():
+    with pytest.raises(AssertionError) as e:
+        labels = data["labels"]
+        pred_probs = data["pred_probs"]
+
+        # baseline scenario where all the pred_probs are the same in the ensemble list
+        num_repeat = 3
+        pred_probs_list = np.repeat(
+            [pred_probs], num_repeat, axis=0
+        )  # this should be a list not an array
+
+        # AssertionError because pred_probs_list is an array
+        _ = rank.get_label_quality_ensemble_scores(labels, pred_probs_list)
+
+        # AssertionError because pred_probs_list is empty
+        _ = rank.get_label_quality_ensemble_scores(labels=labels, pred_probs_list=[])
+
+
+def test_unsupported_method_for_adjust_pred_probs():
+    with pytest.raises(ValueError) as e:
+        labels = data["labels"]
+        pred_probs = data["pred_probs"]
+
+        # method that do not support adjust_pred_probs
+        # note: use a list of methods if there are multiple methods that do not support adjust_pred_probs
+        method = "confidence_weighted_entropy"
+
+        _ = rank.get_label_quality_scores(labels, pred_probs, adjust_pred_probs=True, method=method)
+
+
+def test_find_top_issues():
+    DEFAULT_TOP = 10  # CHANGE THIS IS THE DEFAULT CHANGES
+    X_train = data["X_train"]
+    X_test = data["X_test"]
+    X_ood = np.array([[999999999.0, 999999999.0]])  # Create OOD datapoint
+    X_test_with_ood = np.vstack([X_test, X_ood])  # Add OOD datapoint to X_test
+
+    # Create OOD object (use knn without cosine metric to identify X_ood correctly)
+    knn = NearestNeighbors(n_neighbors=5).fit(X_train)
+    ood_outlier = outlier.OutOfDistribution(params={"knn": knn})
+    ood_scores = ood_outlier.score(features=X_test_with_ood)
+
+    # Get top ood score for outlier example
+    top_outlier_indices = rank.find_top_issues(quality_scores=ood_scores, top=len(ood_scores))
+    top_outlier_indices_more_k = rank.find_top_issues(quality_scores=ood_scores, top=100000)
+
+    ### Check top scores are calculated correctly
+
+    # Checking that X_ood has the smallest outlier score among all the datapoints and outlier scores identifies that
+    assert np.argmin(ood_scores) == (ood_scores.shape[0] - 1)
+    assert len(top_outlier_indices) == len(ood_scores)
+    assert top_outlier_indices[0] == np.argmin(ood_scores)
+
+    # Checking k > len(ood_scores) is same as sorted list of indices
+    assert len(top_outlier_indices) == len(top_outlier_indices_more_k)
+    assert (top_outlier_indices == top_outlier_indices_more_k).all()
+
+    # Get k = DEFAULT_TOP ood scores
+    top_outlier_indices = rank.find_top_issues(ood_scores)
+    assert len(top_outlier_indices) == DEFAULT_TOP
+
+    # Get k < len(ood_scores) ood scores
+    # Assert top k scores are consistent with different length scores vectors
+    for k in [0, 1, 3]:
+        top_outlier_indices_k = rank.find_top_issues(quality_scores=ood_scores, top=k)
+        assert len(top_outlier_indices_k) == k
+        assert (top_outlier_indices_k == top_outlier_indices[:k]).all()  # scores consistent

tests/test_regression.py

@@ -0,0 +1,267 @@
+import pytest
+import random
+import numpy as np
+import pandas as pd
+from sklearn.svm import SVR
+from sklearn.metrics import r2_score
+from sklearn.linear_model import LinearRegression
+from cleanlab.regression.rank import (
+    get_label_quality_scores,
+    _get_residual_score_for_each_label,
+    _get_outre_score_for_each_label,
+)
+from cleanlab.regression.learn import CleanLearning
+
+# set seed for reproducability
+SEED = 1
+np.random.seed(SEED)
+random.seed(SEED)
+
+
+def make_data(num_examples=200, num_features=3, noise=0.2, error_frac=0.1, error_noise=5):
+    X = np.random.random(size=(num_examples, num_features))
+    coefficients = np.random.uniform(-1, 1, size=num_features)
+    label_noise = np.random.normal(scale=noise, size=num_examples)
+
+    true_y = np.dot(X, coefficients)
+    y = np.dot(X, coefficients) + label_noise
+
+    # add extra noisy examples
+    num_errors = int(num_examples * error_frac)
+    extra_noise = np.random.normal(scale=error_noise, size=num_errors)
+    random_idx = np.random.choice(num_examples, num_errors)
+    y[random_idx] += extra_noise
+    error_idx = np.argsort(abs(y - true_y))[-num_errors:]  # get the noisiest examples idx
+
+    # create test set
+    X_test = np.random.random(size=(num_examples, num_features))
+    label_noise = np.random.normal(scale=noise, size=num_examples)
+    y_test = np.dot(X_test, coefficients) + label_noise
+
+    return {
+        "X": X,
+        "y": y,
+        "true_y": true_y,
+        "X_test": X_test,
+        "y_test": y_test,
+        "error_idx": error_idx,
+    }
+
+
+# To be used for most tests
+data = make_data()
+X, labels, predictions = data["X"], data["y"], data["true_y"]
+error_idx = data["error_idx"]
+X_test, y_test = data["X_test"], data["y_test"]
+y = labels  # for ease
+
+# Used for characterization tests
+small_labels = np.array([1, 2, 3, 4])
+small_predictions = np.array([2, 2, 5, 4.1])
+expected_score_outre = np.array([0.2162406, 0.62585509, 0.20275104, 0.62585509])
+expected_score_residual = np.array([0.36787944, 1.0, 0.13533528, 0.90483742])
+expected_scores = {"outre": expected_score_outre, "residual": expected_score_residual}
+
+# Inputs that are not array like
+aConstant = 1
+aString = "predictions_non_array"
+aDict = {"labels": [1, 2], "predictions": [2, 3]}
+aSet = {1, 2, 3, 4}
+aBool = True
+
+
+@pytest.fixture
+def non_array_input():
+    return [aConstant, aString, aDict, aSet, aBool]
+
+
+# test with deafault parameters
+def test_output_shape_type():
+    scores = get_label_quality_scores(labels=labels, predictions=predictions)
+    assert labels.shape == scores.shape
+    assert isinstance(scores, np.ndarray)
+
+
+def test_labels_are_arraylike(non_array_input):
+    for new_input in non_array_input:
+        with pytest.raises(ValueError) as error:
+            get_label_quality_scores(labels=new_input, predictions=predictions)
+            assert error.type == ValueError
+
+
+def test_predictionns_are_arraylike(non_array_input):
+    for new_input in non_array_input:
+        with pytest.raises(ValueError) as error:
+            get_label_quality_scores(labels=labels, predictions=new_input)
+            assert error.type == ValueError
+
+
+# test for input shapes
+def test_input_shape_labels():
+    with pytest.raises(AssertionError) as error:
+        get_label_quality_scores(labels=labels[:-1], predictions=predictions)
+    assert (
+        str(error.value)
+        == f"Number of examples in labels {labels[:-1].shape} and predictions {predictions.shape} are not same."
+    )
+
+
+def test_input_shape_predictions():
+    with pytest.raises(AssertionError) as error:
+        get_label_quality_scores(labels=labels, predictions=predictions[:-1])
+    assert (
+        str(error.value)
+        == f"Number of examples in labels {labels.shape} and predictions {predictions[:-1].shape} are not same."
+    )
+
+
+# test individual scoring functions
+@pytest.mark.parametrize(
+    "scoring_funcs",
+    [_get_residual_score_for_each_label, _get_outre_score_for_each_label],
+)
+def test_individual_scoring_functions(scoring_funcs):
+    scores = scoring_funcs(labels=labels, predictions=predictions)
+    assert labels.shape == scores.shape
+    assert isinstance(scores, np.ndarray)
+
+
+# test for method argument
+@pytest.mark.parametrize(
+    "method",
+    [
+        "residual",
+        "outre",
+    ],
+)
+def test_method_pass_get_label_quality_scores(method):
+    scores = get_label_quality_scores(labels=labels, predictions=predictions, method=method)
+    assert labels.shape == scores.shape
+    assert isinstance(scores, np.ndarray)
+
+
+@pytest.mark.parametrize(
+    "method",
+    [
+        "residual",
+        "outre",
+    ],
+)
+def test_expected_scores(method):
+    # characterization test
+    scores = get_label_quality_scores(
+        labels=small_labels, predictions=small_predictions, method=method
+    )
+    assert np.allclose(scores, expected_scores[method], atol=1e-08)
+
+
+def test_cleanlearning():
+    # test fit and predict
+    cl = CleanLearning()
+    cl.fit(X, y)
+    preds = cl.predict(X)
+    cl_r2_score = cl.score(X, y)
+    manual_r2_score = r2_score(y, preds)
+    assert len(preds) == len(y)
+    assert isinstance(cl_r2_score, float)
+    assert cl_r2_score == manual_r2_score
+
+    # check if label issues were identified
+    label_issues = cl.get_label_issues()
+    identified_label_issues = label_issues[label_issues["is_label_issue"] == True].index
+    frac_errors_identified = np.mean([e in identified_label_issues for e in error_idx])
+    assert frac_errors_identified >= 0.9  # assert most errors were detected
+
+    # compare perf to base LinearRegression model
+    cl_score = cl.score(X_test, y_test)
+    lr = LinearRegression()
+    lr.fit(X, y)
+    lr_score = lr.score(X_test, y_test)
+    assert cl_score > lr_score
+
+    # test passing in label issues in various forms
+    # also test different regression model
+    cl = CleanLearning(model=SVR())
+    label_issues = cl.find_label_issues(X, y)
+    assert isinstance(label_issues, pd.DataFrame)
+
+    cl.fit(X, y, label_issues=label_issues)
+    cl.fit(X, pd.Series(y), label_issues=label_issues["is_label_issue"])
+    cl.fit(X, list(y), label_issues=label_issues["is_label_issue"].values)
+
+
+def test_optional_inputs():
+    # test with sample_weight input
+    cl = CleanLearning(verbose=1)
+    cl.fit(X, y, sample_weight=np.random.random(size=len(y)))
+    cl.fit(X, y, label_issues=cl.get_label_issues(), sample_weight=np.random.random(size=len(y)))
+
+    # test with uncertainty input
+    cl = CleanLearning()
+    cl.find_label_issues(X, y, uncertainty=5)  # constant uncertainty
+    cl.find_label_issues(X, y, uncertainty=np.random.random(size=len(y)))  # per-example uncertainty
+
+    # test with not calculating uncertainty
+    cl = CleanLearning(n_boot=0, include_aleatoric_uncertainty=False)
+    cl.find_label_issues(X, y)
+
+    # test with odd grid search sizes
+    cl = CleanLearning()
+    cl.find_label_issues(X, y, coarse_search_range=[0.2])
+    cl.find_label_issues(X, y, fine_search_size=0)
+    cl.fit(
+        X, y, find_label_issues_kwargs={"coarse_search_range": [0.2, 0.1], "fine_search_size": 2}
+    )
+
+
+def test_low_example_count():
+    data_tiny = make_data(num_examples=3)
+    X_tiny, y_tiny = data_tiny["X"], data_tiny["y"]
+
+    try:
+        cl = CleanLearning()
+        cl.find_label_issues(X_tiny, y_tiny)
+    except ValueError as e:
+        assert "There are too few examples" in str(e)
+
+    cl = CleanLearning(cv_n_folds=3)
+    cl.find_label_issues(X_tiny, y_tiny)
+    assert isinstance(cl.get_label_issues(), pd.DataFrame)
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_save_space():
+    # test label issues df does not save
+    cl = CleanLearning()
+    cl.find_label_issues(X, y, save_space=True)
+    assert cl.get_label_issues() is None
+
+    # test label issues df deletes properly
+    cl = CleanLearning()
+    cl.find_label_issues(X, y)
+    assert isinstance(cl.get_label_issues(), pd.DataFrame)
+
+    cl.save_space()
+    assert cl.get_label_issues() is None
+
+
+@pytest.mark.parametrize("N", [10, 100, 1000])
+@pytest.mark.parametrize("method", ["residual", "outre"])
+def test_all_identical_examples(N, method):
+
+    # All examples have predictions identical to the given labels/targets
+    labels = np.zeros(N)
+    predictions = np.copy(labels)
+
+    # Except the last ~quarter of examples have labels that are further away from the predictions
+    cutoff_index = N // 4
+    predictions[-cutoff_index:] += 1
+
+    scores = get_label_quality_scores(labels=labels, predictions=predictions, method=method)
+    np.testing.assert_allclose(scores[:-cutoff_index], 1, atol=1e-04)
+    if method == "outre":
+        # Assert that the scores for the last (bad) ~quarter of examples are close to 0
+        np.testing.assert_allclose(scores[-cutoff_index:], 0, atol=1e-04)
+    else:
+        # Residual method should give "imperfect" scores for the last ~quarter of examples, but not necessarily near 0
+        assert np.all(scores[-cutoff_index:] < 1)

tests/test_segmentation.py

@@ -0,0 +1,486 @@
+"""
+Scripts to test cleanlab.segmentation package
+"""
+
+import numpy as np
+import random
+
+np.random.seed(0)
+import pytest
+from unittest import mock
+import matplotlib
+
+matplotlib.use("Agg")  # Set non-interactive backend before importing pyplot
+import matplotlib.pyplot as plt
+from pathlib import Path
+
+from cleanlab.internal.multilabel_scorer import softmin
+
+# Segmentation utils
+from cleanlab.internal.segmentation_utils import (
+    _check_input,
+    _get_valid_optional_params,
+    _get_summary_optional_params,
+)
+
+# Filter
+from cleanlab.segmentation.filter import (
+    find_label_issues,
+)
+
+# Rank
+from cleanlab.segmentation.rank import (
+    get_label_quality_scores,
+    issues_from_scores,
+    _get_label_quality_per_image,
+)
+
+# Summary
+from cleanlab.segmentation.summary import (
+    display_issues,
+    common_label_issues,
+    filter_by_class,
+    _generate_colormap,
+)
+
+
+def generate_three_image_dataset(bad_index):
+    good_gt = np.zeros((10, 10))
+    good_gt[:5, :] = 1.0
+    bad_gt = np.ones((10, 10))
+    bad_gt[:5, :] = 0.0
+    good_pr = np.random.random((2, 10, 10))
+    good_pr[0, :5, :] = good_pr[0, :5, :] / 10
+    good_pr[1, 5:, :] = good_pr[1, 5:, :] / 10
+
+    val = np.binary_repr([4, 2, 1][bad_index], width=3)
+    error = [int(case) for case in val]
+
+    labels = []
+    pred = []
+    for case in val:
+        if case == "0":
+            labels.append(good_gt)
+            pred.append(good_pr)
+        else:
+            labels.append(bad_gt)
+            pred.append(good_pr)
+
+    labels = np.array(labels)
+    pred_probs = np.array(pred)
+    return labels, pred_probs, error
+
+
+labels, pred_probs, error = generate_three_image_dataset(random.randint(0, 2))
+labels, pred_probs = labels.astype(int), pred_probs.astype(float)
+num_images, num_classes, h, w = pred_probs.shape
+
+
+def test_find_label_issues():
+    issues = find_label_issues(labels, pred_probs, n_jobs=None, batch_size=1000)
+    assert np.argmax(error) == np.argmax(issues.sum((1, 2)))
+
+    issues = find_label_issues(labels, pred_probs, downsample=2, batch_size=1739)
+    assert np.argmax(error) == np.argmax(issues.sum((1, 2)))
+
+    issues = find_label_issues(labels, pred_probs, downsample=5, n_jobs=None, batch_size=2838)
+    assert np.argmax(error) == np.argmax(issues.sum((1, 2)))
+
+    with pytest.raises(Exception) as e:
+        issues = find_label_issues(labels, pred_probs, downsample=4, n_jobs=None, batch_size=1000)
+
+    # Simple tests
+    # Test case 1: Test with larger batch_size
+    issues = find_label_issues(labels, pred_probs, downsample=1, n_jobs=None, batch_size=2000)
+    assert np.argmax(error) == np.argmax(issues.sum((1, 2)))
+
+    # Test case 2: Test with smaller batch_size
+    issues = find_label_issues(labels, pred_probs, downsample=1, n_jobs=None, batch_size=500)
+    assert np.argmax(error) == np.argmax(issues.sum((1, 2)))
+
+    # Test case 3: Test verbose off
+    issues = find_label_issues(labels, pred_probs, downsample=1, verbose=False)
+
+    assert np.argmax(error) == np.argmax(issues.sum((1, 2)))
+
+    # Test case 5: Test with invalid downsample value
+    with pytest.raises(Exception) as e:
+        issues = find_label_issues(labels, pred_probs, downsample=3, n_jobs=None, batch_size=1000)
+
+    # Test case 6: Test with n_jobs parameter
+    issues = find_label_issues(labels, pred_probs, downsample=1, n_jobs=2, batch_size=1000)
+    assert np.argmax(error) == np.argmax(issues.sum((1, 2)))
+
+    # Test case 7: Test with invalid labels
+    with pytest.raises(Exception) as e:
+        issues = find_label_issues(
+            np.array([[[[1, 2, 3]]]]), pred_probs, downsample=1, n_jobs=None, batch_size=1000
+        )
+
+    # Test case 8: Test with invalid pred_probs
+    with pytest.raises(Exception) as e:
+        issues = find_label_issues(
+            labels, np.array([[[[0.1, 0.2, 0.3]]]]), downsample=1, n_jobs=None, batch_size=1000
+        )
+
+
+def test_results_are_consistent_with_batch_size(tmp_path: Path):
+    """
+    Test that find_label_issues works with large memmap arrays and different batch sizes
+    """
+
+    # Create dummy versions of pred_probs and labels
+    # write to the pytest tmp_path so that the files are deleted after the test
+    pred_probs_file = tmp_path / "pred_probs.npy"
+    labels_file = tmp_path / "labels.npy"
+    np.save(pred_probs_file, np.random.rand(100, 2, 5, 5))
+    np.save(labels_file, np.random.randint(0, 2, (100, 5, 5)))
+
+    # Load the numpy arrays from disk
+    pred_probs = np.load(pred_probs_file, mmap_mode="r")
+    pred_labels = np.load(labels_file, mmap_mode="r")
+
+    # Test with different batch sizes
+    batch_sizes = [1, 50, 100]
+    issues_list = []
+    for batch_size in batch_sizes:
+        issues = find_label_issues(pred_labels, pred_probs, n_jobs=None, batch_size=batch_size)
+        issues_list.append(issues)
+
+    # Verify that the results are identical regardless of the batch size
+    for i in range(len(batch_sizes) - 1):
+        assert np.array_equal(issues_list[i], issues_list[i + 1])
+
+
+def test_find_label_issues_sizes():
+    # checks inputs of different sizes
+    labels, pred_probs = np.random.randint(0, 2, (2, 9, 7)), np.random.random((2, 2, 9, 7))
+    issues = find_label_issues(labels, pred_probs)
+
+    labels, pred_probs = np.random.randint(0, 2, (2, 13, 47)), np.random.random((2, 2, 13, 47))
+    issues = find_label_issues(labels, pred_probs)
+
+    for _ in range(5):
+        h, w = np.random.randint(1, 100, 2)
+        labels, pred_probs = np.random.randint(0, 2, (2, h, w)), np.random.random((2, 2, h, w))
+        issues = find_label_issues(labels, pred_probs)
+
+
+def test__check_input():
+    bad_gt = np.random.random((5, 10, 20))
+    with pytest.raises(Exception) as e:
+        _check_input(bad_gt, bad_gt)
+
+    bad_pr = np.random.random((5, 2, 10, 20))
+    with pytest.raises(Exception) as e:
+        _check_input(bad_pr, bad_pr)
+
+    smaller_pr = np.random.random((5, 2, 9, 20))
+    with pytest.raises(Exception) as e:
+        _check_input(bad_gt, smaller_pr)
+
+    fewer_gt = np.random.random((4, 10, 20))
+    with pytest.raises(Exception) as e:
+        _check_input(fewer_gt, smaller_pr)
+
+
+@pytest.mark.filterwarnings("ignore::UserWarning")  # Should be 1 warning
+def test_get_label_quality_scores():
+    image_scores_softmin, pixel_scores = get_label_quality_scores(
+        labels, pred_probs, method="softmin"
+    )
+    assert np.argmax(error) == np.argmin(image_scores_softmin)
+
+    image_scores_npi, pixel_scores = get_label_quality_scores(
+        labels, pred_probs, method="num_pixel_issues", downsample=1
+    )
+    assert np.argmax(error) == np.argmin(image_scores_npi)
+
+    image_scores_softmin, pixel_scores = get_label_quality_scores(
+        labels, pred_probs, downsample=1, method="softmin"
+    )
+    assert len(image_scores_softmin) == labels.shape[0]
+    assert pixel_scores.shape == labels.shape
+
+
+@pytest.mark.parametrize(
+    "method, downsample, batch_size, expected_exception, expected_message",
+    [
+        (
+            "num_pixel_issues",
+            4,
+            None,
+            Exception,
+            "Height 10 and width 10 not divisible by downsample value of 4. Set kwarg downsample to 1 to avoid downsampling.",
+        ),
+        (
+            "invalid_method",
+            None,
+            None,
+            Exception,
+            "Invalid Method: Specify correct method. Currently only supports 'softmin'",
+        ),
+        ("num_pixel_issues", 1, -1, ValueError, "Batch size must be greater than 0, got -1"),
+        ("num_pixel_issues", 1, 0, ValueError, "Batch size must be greater than 0, got 0"),
+    ],
+    ids=["downsample", "method", "batch_size_negative", "batch_size_zero"],
+)
+def test_get_label_quality_scores_exceptions(
+    method, downsample, batch_size, expected_exception, expected_message
+):
+    args = {"labels": labels, "pred_probs": pred_probs, "method": method}
+    if downsample is not None:
+        args["downsample"] = downsample
+    if batch_size is not None:
+        args["batch_size"] = batch_size
+
+    with pytest.raises(expected_exception) as exc_info:
+        get_label_quality_scores(**args)
+
+    assert expected_message in str(exc_info.value)
+
+
+# different size inpits
+def test_get_label_quality_scores_sizes():
+    # checks inputs of different sizes
+    labels, pred_probs = np.random.randint(0, 2, (2, 9, 7)), np.random.random((2, 2, 9, 7))
+    image_scores_softmin, pixel_scores = get_label_quality_scores(
+        labels, pred_probs, method="softmin"
+    )
+
+    labels, pred_probs = np.random.randint(0, 2, (2, 13, 47)), np.random.random((2, 2, 13, 47))
+    image_scores_softmin, pixel_scores = get_label_quality_scores(
+        labels, pred_probs, method="softmin"
+    )
+
+    for _ in range(5):
+        h, w = np.random.randint(1, 100, 2)
+        labels, pred_probs = np.random.randint(0, 2, (2, h, w)), np.random.random((2, 2, h, w))
+        image_scores_softmin, pixel_scores = get_label_quality_scores(
+            labels, pred_probs, method="softmin"
+        )
+
+
+# Testing issues from scores
+def test_issues_from_scores():
+    image_scores_softmin, pixel_scores = get_label_quality_scores(
+        labels, pred_probs, method="softmin"
+    )
+    issues_from_score = issues_from_scores(image_scores_softmin, pixel_scores, threshold=1)
+    assert np.shape(issues_from_score) == np.shape(pixel_scores)
+    assert h * w * num_images == issues_from_score.sum()
+
+    issues_from_score = issues_from_scores(image_scores_softmin, pixel_scores, threshold=0)
+    assert 0 == issues_from_score.sum()
+
+    issues_from_score = issues_from_scores(image_scores_softmin, pixel_scores, threshold=0.5)
+    assert np.argmax(error) == np.argmax(issues_from_score.sum((1, 2)))
+
+    sort_by_score = issues_from_scores(image_scores_softmin, threshold=0.5)
+    assert error[sort_by_score[0]] == 1
+
+
+def test_issues_from_scores_no_pixel_scores():
+    # Test if function works correctly when pixel_scores is None
+    image_scores_softmin, _ = get_label_quality_scores(labels, pred_probs, method="softmin")
+    issues_from_score_result = issues_from_scores(image_scores_softmin, None, threshold=1)
+    assert np.shape(issues_from_score_result) == (num_images,)
+
+
+def test_issues_from_scores_various_thresholds():
+    # Test if function works correctly for various values of threshold
+    image_scores_softmin, pixel_scores = get_label_quality_scores(
+        labels, pred_probs, method="softmin"
+    )
+    for threshold in [0.1, 0.5, 0.9]:
+        issues_from_score_result = issues_from_scores(
+            image_scores_softmin, pixel_scores, threshold=threshold
+        )
+        assert np.all(issues_from_score_result == (pixel_scores < threshold))
+
+
+def test_issues_from_scores_invalid_inputs():
+    # Test if function raises exception when input parameters are invalid
+    with pytest.raises(ValueError):
+        issues_from_scores(None)
+    with pytest.raises(ValueError):
+        issues_from_scores(np.array([0.1, 0.2, 0.3]), threshold=1.1)  # Threshold more than 1
+    with pytest.raises(ValueError):
+        issues_from_scores(np.array([0.1, 0.2, 0.3]), threshold=-0.1)  # Threshold less than 0
+
+
+def test_issues_from_scores_different_input_sizes():
+    # Test if function works correctly for different sizes of input arrays
+    for num_images in range(1, 5):
+        image_scores = np.random.rand(num_images)
+        pixel_scores = np.random.rand(num_images, 100, 100)
+        issues_from_score_result = issues_from_scores(image_scores, pixel_scores, threshold=0.5)
+        assert np.shape(issues_from_score_result) == np.shape(pixel_scores)
+
+
+def test_issues_from_scores_sorting():
+    # Test if function correctly sorts image_scores
+    image_scores_softmin, _ = get_label_quality_scores(labels, pred_probs, method="softmin")
+    issues_from_score_result = issues_from_scores(image_scores_softmin, None, threshold=0.5)
+    assert np.all(np.sort(image_scores_softmin) == image_scores_softmin[issues_from_score_result])
+
+
+def test__get_label_quality_per_image():
+    # Test when pixel_scores is a random list of 100 values, method is "softmin", and temperature is random
+    random_score_array = np.random.random((100,))
+    temp = random.random()
+    score = _get_label_quality_per_image(random_score_array, method="softmin", temperature=temp)
+
+    cleanlab_softmin = softmin(
+        np.expand_dims(random_score_array, axis=0), axis=1, temperature=temp
+    )[0]
+    assert cleanlab_softmin == score, "Expected cleanlab_softmin to be equal to score"
+
+    # Test when pixel_scores is an empty list, should raise an error
+    empty_score_array = np.array([])
+    with pytest.raises(Exception) as e:
+        _get_label_quality_per_image(empty_score_array, method="softmin", temperature=temp)
+
+    # Test when method is None
+    with pytest.raises(Exception):
+        _get_label_quality_per_image(random_score_array, method=None, temperature=temp)
+
+    # Test when method is not "softmin", should raise an exception
+    with pytest.raises(Exception):
+        _get_label_quality_per_image(random_score_array, method="invalid_method", temperature=temp)
+
+    #     Test when temperature is 0, should raise an error
+    with pytest.raises(Exception):
+        _get_label_quality_per_image(random_score_array, method="softmin", temperature=0)
+
+    with pytest.raises(Exception):
+        _get_label_quality_per_image(random_score_array, method="softmin", temperature=None)
+
+
+def test_generate_color_map():
+    colors = _generate_colormap(0)
+    assert len(colors) == 0
+
+    colors = _generate_colormap(1)
+    assert len(colors) == 1
+    assert len(colors[0]) == 4
+
+    colors = _generate_colormap(-1)
+    assert len(colors) == 0
+
+    colors = _generate_colormap(5)
+    assert len(colors) == 5
+
+    # test unique
+    num_colors = 385  # max number of unique colors avalible
+    colors = _generate_colormap(num_colors)
+    unique_rows = np.unique(colors, axis=0)
+    assert unique_rows.shape[0] == num_colors
+
+
+def test_display_issues(monkeypatch):
+    monkeypatch.setattr(plt, "show", lambda: None)
+    issues = find_label_issues(labels, pred_probs, downsample=1, n_jobs=None, batch_size=1000)
+    display_issues(issues, top=1)
+
+    display_issues(issues, pred_probs=pred_probs, labels=labels, top=2, class_names=["one", "two"])
+
+    display_issues(issues, pred_probs=pred_probs, labels=labels, top=2)
+    display_issues(issues, labels=labels, top=2)
+    display_issues(issues, pred_probs=pred_probs, top=2)
+
+    # too many issues for top
+    display_issues(issues, pred_probs=pred_probs, labels=labels, top=len(issues) + 5)
+
+    class_issues = filter_by_class(0, issues, labels=labels, pred_probs=pred_probs)
+    display_issues(class_issues, pred_probs=pred_probs, labels=labels, top=2)
+
+    image_scores, pixel_scores = image_scores_softmin, pixel_scores = get_label_quality_scores(
+        labels, pred_probs, method="softmin"
+    )
+    issues_from_score = issues_from_scores(image_scores, pixel_scores, threshold=0.5)
+    display_issues(issues_from_score, pred_probs=pred_probs, labels=labels, top=2)
+
+    display_issues(issues_from_score, pred_probs=pred_probs, labels=labels, top=2, exclude=[0])
+
+    with pytest.raises(ValueError) as e:
+        display_issues(issues_from_score, pred_probs=pred_probs, labels=None, top=2, exclude=[0])
+
+
+@mock.patch("matplotlib.pyplot.figure")
+def test_display_issues_figure(mock_plt, monkeypatch):
+    monkeypatch.setattr(plt, "show", lambda: None)
+    issues = find_label_issues(labels, pred_probs, downsample=1, n_jobs=None, batch_size=1000)
+    display_issues(issues, pred_probs=pred_probs, labels=labels, top=2, class_names=["one", "two"])
+    assert mock_plt.called
+
+
+@mock.patch("matplotlib.pyplot.show")
+def test_display_issues_show(mock_plt):
+    issues = find_label_issues(labels, pred_probs, downsample=1, n_jobs=None, batch_size=1000)
+    display_issues(issues, top=1)
+    assert mock_plt.called
+
+
+def test_common_label_issues(capsys):
+    issues = find_label_issues(labels, pred_probs, downsample=1, n_jobs=None, batch_size=1000)
+
+    # test exclude
+    df = common_label_issues(issues, labels, pred_probs)
+    df_without_0 = common_label_issues(issues, labels, pred_probs, exclude=[0])
+    assert df.shape != df_without_0.shape
+
+    # test verbose
+    captured_words = capsys.readouterr()
+    df = common_label_issues(issues, labels, pred_probs, verbose=False)
+    captured_no_words = capsys.readouterr()
+    assert len(captured_no_words.out) == 0
+    assert len(captured_words.out) > 0
+
+    # test class names & top
+    df_class_names = common_label_issues(issues, labels, pred_probs, class_names=["one", "two"])
+    captured_top_all = capsys.readouterr()
+    df_top_1 = common_label_issues(issues, labels, pred_probs, top=1)
+    captured_top_1 = capsys.readouterr()
+    assert len(captured_top_1.out) < len(captured_top_all.out)
+    assert df_class_names["given_label"].to_list() != df["given_label"].to_list()
+
+
+def test_filter_by_class():
+    issues = find_label_issues(labels, pred_probs, downsample=1, n_jobs=None, batch_size=1000)
+    class_0_issues = filter_by_class(0, issues, labels=labels, pred_probs=pred_probs)
+    class_1_issues = filter_by_class(1, issues, labels=labels, pred_probs=pred_probs)
+
+    class_300_issues = filter_by_class(300, issues, labels=labels, pred_probs=pred_probs)
+
+    assert (class_0_issues == class_1_issues).all()  # mirror issues
+    assert np.sum(class_300_issues) == 0
+
+
+def test_summary_sizes(monkeypatch):
+    monkeypatch.setattr(plt, "show", lambda: None)
+
+    for _ in range(5):
+        h, w = np.random.randint(1, 100, 2)
+        labels, pred_probs = np.random.randint(0, 2, (2, h, w)), np.random.random((2, 2, h, w))
+        issues = find_label_issues(labels, pred_probs, downsample=1, n_jobs=None, batch_size=1000)
+        class_300_issues = filter_by_class(0, issues, labels=labels, pred_probs=pred_probs)
+        df = common_label_issues(issues, labels, pred_probs)
+        display_issues(issues)
+
+
+def test_get_valid_functions():
+    optional_batch_size = 10
+    optional_n_jobs = 2
+    x, y = _get_valid_optional_params(optional_batch_size, optional_n_jobs)
+    assert x == optional_batch_size and y == optional_n_jobs
+    x, y = _get_valid_optional_params(None, None)
+    assert x == 10000 and y == None
+
+    optional_class_names = [1, 2]
+    optional_exclude = [1]
+    optional_top = 10
+    x, y, z = _get_summary_optional_params(optional_class_names, optional_exclude, optional_top)
+    assert x == optional_class_names and y == optional_exclude and z == optional_top
+    x, y, z = _get_summary_optional_params(None, None, None)
+    assert x == None and y == [] and z == 20

tests/test_span_classification.py

@@ -0,0 +1,57 @@
+from cleanlab.experimental.span_classification import (
+    find_label_issues,
+    display_issues,
+    get_label_quality_scores,
+)
+import numpy as np
+import pytest
+import warnings
+
+warnings.filterwarnings("ignore")
+words = [["I", "love", "Cleanlab", "Inc"], ["A", "new", "park"]]
+
+pred_probs = [
+    np.array([0.3, 0.2, 0.9, 0.1]),
+    np.array([0.1, 0.1, 0.9]),
+]
+
+labels = [[0, 0, 1, 1], [0, 0, 1]]
+class_names = ["O", "Span"]
+
+
+@pytest.mark.parametrize(
+    "test_labels",
+    [labels, [np.array(l) for l in labels]],
+    ids=["list labels", "np.array labels"],
+)
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_find_label_issues(test_labels):
+    issues = find_label_issues(test_labels, pred_probs)
+    assert isinstance(issues, list)
+    assert len(issues) == 1
+    assert issues[0] == (0, 3)
+
+
+issues = find_label_issues(labels, pred_probs)
+
+
+def test_display_issues():
+    display_issues(issues, words)
+    display_issues(issues, tokens=words, labels=labels)
+    display_issues(issues, words, pred_probs=pred_probs)
+    display_issues(issues, words, pred_probs=pred_probs, labels=labels)
+    display_issues(issues, words, pred_probs=pred_probs, labels=labels, class_names=class_names)
+
+
+@pytest.fixture(name="label_quality_scores")
+def fixture_label_quality_scores():
+    sentence_scores, token_info = get_label_quality_scores(labels, pred_probs)
+    return sentence_scores, token_info
+
+
+def test_get_label_quality_scores(label_quality_scores):
+    sentence_scores, token_info = label_quality_scores
+    assert len(sentence_scores) == 2
+    assert np.allclose(sentence_scores, [0.1, 0.9])
+    assert len(token_info) == 2
+    assert np.allclose(token_info[0], [0.7, 0.8, 0.9, 0.1])

tests/test_token_classification.py

@@ -0,0 +1,343 @@
+from cleanlab.internal.token_classification_utils import (
+    get_sentence,
+    filter_sentence,
+    process_token,
+    mapping,
+    merge_probs,
+    color_sentence,
+    _replace_sentence,
+)
+from cleanlab.token_classification.filter import find_label_issues
+from cleanlab.token_classification.rank import (
+    get_label_quality_scores,
+    issues_from_scores,
+    _softmin_sentence_score,
+)
+from cleanlab.token_classification.summary import (
+    display_issues,
+    common_label_issues,
+    filter_by_token,
+)
+import numpy as np
+import pandas as pd
+import pytest
+
+import warnings
+
+warnings.filterwarnings("ignore")
+words = [["Hello", "World"], ["#I", "love", "Cleanlab"], ["A"]]
+sentences = ["Hello World", "#I love Cleanlab", "A"]
+
+pred_probs = [
+    np.array([[0.9, 0.1, 0], [0.6, 0.2, 0.2]]),
+    np.array([[0.1, 0, 0.9], [0.1, 0.8, 0.1], [0.1, 0.8, 0.1]]),
+    np.array([[0.1, 0.1, 0.8]]),
+]
+
+
+labels = [[0, 0], [1, 1, 1], [2]]
+
+maps = [0, 1, 0, 1]
+class_names = ["A", "B", "C", "D"]
+
+
+def test_get_sentence():
+    actual_sentences = list(map(get_sentence, words))
+    assert actual_sentences == sentences
+
+    # Test with allowed special characters
+    words_separated_by_hyphen = ["Heading", "-", "Title"]
+    assert get_sentence(words_separated_by_hyphen) == "Heading - Title"
+
+    words_within_parentheses = ["Some", "reason", "(", "Explanation", ")"]
+    assert get_sentence(words_within_parentheses) == "Some reason (Explanation)"
+
+
+def test_filter_sentence():
+    filtered_sentences, mask = filter_sentence(sentences)
+    assert filtered_sentences == ["Hello World"]
+    assert mask == [True, False, False]
+
+    filtered_sentences, mask = filter_sentence(sentences, lambda x: len(x) > 1)
+    assert filtered_sentences == ["Hello World", "#I love Cleanlab"]
+    assert mask == [True, True, False]
+
+    filtered_sentences, mask = filter_sentence(sentences, lambda x: "#" not in x)
+    assert filtered_sentences == ["Hello World", "A"]
+    assert mask == [True, False, True]
+
+
+def test_process_token():
+    test_cases = [
+        ("Cleanlab", [("C", "a")], "aleanlab"),
+        ("Cleanlab", [("C", "a"), ("a", "C")], "aleCnlCb"),
+    ]
+    for token, replacements, expected in test_cases:
+        processed = process_token(token, replacements)
+        assert processed == expected
+
+
+def test_mapping():
+    test_cases = [(l, expected) for l, expected in zip(labels, [[0, 0], [1, 1, 1], [0]])]
+    for l, expected in test_cases:
+        mapped = mapping(l, maps)
+        assert mapped == expected
+
+
+def test_merge_probs():
+    merged_probs = merge_probs(pred_probs[0], maps)
+    expected = np.array([[0.9, 0.1], [0.8, 0.2]])
+    assert np.allclose(expected, merged_probs)
+
+    merged_probs = merge_probs(pred_probs[1], maps)
+    expected = np.array([[1.0, 0.0], [0.2, 0.8], [0.2, 0.8]])
+    assert np.allclose(expected, merged_probs)
+
+    merged_probs = merge_probs(pred_probs[2], maps)
+    expected = np.array([[0.9, 0.1]])
+    assert np.allclose(expected, merged_probs)
+
+
+def test_merge_probs_with_normalization():
+    # Ignore probabilities for class/entity 0
+    norm_maps = [-1, 1, 0, 1]
+    merged_probs = merge_probs(pred_probs[0], norm_maps)
+    expected = np.array([[0.0, 1.0], [0.5, 0.5]])
+    assert np.allclose(expected, merged_probs)
+
+    merged_probs = merge_probs(pred_probs[1], norm_maps)
+    expected = np.array([[1.0, 0.0], [1 / 9, 8 / 9], [1 / 9, 8 / 9]])
+    assert np.allclose(expected, merged_probs)
+
+    merged_probs = merge_probs(pred_probs[2], norm_maps)
+    expected = np.array([[8 / 9, 1 / 9]])
+
+    # Ignore probabilities for class/entity 1
+    norm_maps = [0, -1, 0, 1]
+    merged_probs = merge_probs(pred_probs[0], norm_maps)
+    expected = np.array([[1.0, 0.0], [1.0, 0.0]])
+    assert np.allclose(expected, merged_probs)
+
+    merged_probs = merge_probs(pred_probs[1], norm_maps)
+    expected = np.array([[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]])
+    assert np.allclose(expected, merged_probs)
+
+    merged_probs = merge_probs(pred_probs[2], norm_maps)
+    expected = np.array([[1.0, 0.0]])
+    assert np.allclose(expected, merged_probs)
+
+
+# Color boundaries
+C_L, C_R = "\x1b[31m", "\x1b[0m"
+
+
+@pytest.mark.parametrize(
+    "sentence,word,expected",
+    [
+        ("Hello World", "World", f"Hello {C_L}World{C_R}"),
+        ("Hello World", "help", "Hello World"),
+        ("If you and I were to meet", "I", f"If you and {C_L}I{C_R} were to meet"),
+        ("If you and I were to meet", "If you and I", f"{C_L}If you and I{C_R} were to meet"),
+        ("If you and I were to meet", "If you and I w", f"{C_L}If you and I w{C_R}ere to meet"),
+        ("I think I know this", "I", f"{C_L}I{C_R} think {C_L}I{C_R} know this"),
+        ("A good reason for a test", "a", f"A good reason for {C_L}a{C_R} test"),
+        ("ab ab a b ab", "ab a", f"ab {C_L}ab a{C_R} b ab"),
+        ("ab ab ab ab", "ab a", f"{C_L}ab a{C_R}b {C_L}ab a{C_R}b"),
+        (
+            "Alan John Percivale (A.j.p.) Taylor died",
+            "(",
+            f"Alan John Percivale {C_L}({C_R}A.j.p.) Taylor died",
+        ),
+    ],
+    ids=[
+        "single_word",
+        "no_match",
+        "ignore_subwords",
+        "multi-token_match",
+        "substring_replacement",
+        "multiple_matches",
+        "case_sensitive",
+        "only_word_boundary",
+        "non_overlapping_substrings",
+        "issue_403-escape_special_regex_characters",
+    ],
+)
+def test_color_sentence(monkeypatch: pytest.MonkeyPatch, sentence, word, expected):
+    import os
+
+    monkeypatch.setattr(os, "isatty", lambda fd: True)
+    monkeypatch.setattr("sys.stdout.isatty", lambda: True)
+    monkeypatch.setattr("sys.stdout.fileno", lambda: 1)
+
+    colored = color_sentence(sentence, word)
+    assert colored == expected
+
+
+@pytest.mark.parametrize(
+    "sentence,word,expected",
+    [
+        ("Hello World", "World", "Hello [EXPECTED]"),
+        ("Hello World", "help", "Hello World"),
+        ("If you and I were to meet", "I", "If you and [EXPECTED] were to meet"),
+        ("If you and I were to meet", "If you and I", "[EXPECTED] were to meet"),
+        ("If you and I were to meet", "If you and I w", "[EXPECTED]ere to meet"),
+        ("I think I know this", "I", "[EXPECTED] think [EXPECTED] know this"),
+        ("A good reason for a test", "a", "A good reason for [EXPECTED] test"),
+        ("ab ab a b ab", "ab a", "ab [EXPECTED] b ab"),
+        ("ab ab ab ab", "ab a", "[EXPECTED]b [EXPECTED]b"),
+        (
+            "Alan John Percivale (A.j.p.) Taylor died",
+            "(",
+            "Alan John Percivale [EXPECTED]A.j.p.) Taylor died",
+        ),
+    ],
+    ids=[
+        "single_word",
+        "no_match",
+        "ignore_subwords",
+        "multi-token_match",
+        "substring_replacement",
+        "multiple_matches",
+        "case_sensitive",
+        "only_word_boundary",
+        "non_overlapping_substrings",
+        "issue_403-escape_special_regex_characters",
+    ],
+)
+def test_replace_sentence(sentence, word, expected):
+    new_sentence = _replace_sentence(sentence, word, "[EXPECTED]")
+    assert new_sentence == expected
+
+
+issues = find_label_issues(labels, pred_probs)
+
+
+@pytest.mark.parametrize(
+    "test_labels",
+    [labels, [np.array(l) for l in labels]],
+    ids=["list labels", "np.array labels"],
+)
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_find_label_issues(test_labels):
+    issues = find_label_issues(test_labels, pred_probs)
+    assert isinstance(issues, list)
+    assert len(issues) == 1
+    assert issues[0] == (1, 0)
+    issues2 = find_label_issues(
+        test_labels, pred_probs, return_indices_ranked_by="normalized_margin", n_jobs=1
+    )
+    assert isinstance(issues2, list)
+    # Compare results with low_memory=True. Pass unused argument n_jobs=1
+    issues_lm = find_label_issues(test_labels, pred_probs, low_memory=True, n_jobs=1)
+    intersection = len(list(set(issues).intersection(set(issues_lm))))
+    union = len(set(issues)) + len(set(issues_lm)) - intersection
+    assert float(intersection) / union > 0.95
+
+
+def test_softmin_sentence_score():
+    token_scores = [[0.9, 0.6], [0.0, 0.8, 0.8], [0.8]]
+    sentence_scores = _softmin_sentence_score(token_scores)
+    assert isinstance(sentence_scores, np.ndarray)
+    assert np.allclose(sentence_scores, [0.60074, 1.8e-07, 0.8])
+
+    # Temperature limits
+    sentence_scores = _softmin_sentence_score(token_scores, temperature=0)
+    assert np.allclose(sentence_scores, [0.6, 0.0, 0.8])
+
+    sentence_scores = _softmin_sentence_score(token_scores, temperature=np.inf)
+    assert np.allclose(sentence_scores, [0.75, 1.6 / 3, 0.8])
+
+
+@pytest.fixture(name="label_quality_scores")
+def fixture_label_quality_scores():
+    sentence_scores, token_info = get_label_quality_scores(labels, pred_probs)
+    return sentence_scores, token_info
+
+
+def test_get_label_quality_scores(label_quality_scores):
+    sentence_scores, token_info = label_quality_scores
+    assert len(sentence_scores) == 3
+    assert np.allclose(sentence_scores, [0.6, 0, 0.8])
+    assert len(token_info) == 3
+    assert np.allclose(token_info[0], [0.9, 0.6])
+    sentence_scores_softmin, _ = get_label_quality_scores(
+        labels, pred_probs, sentence_score_method="softmin", tokens=words
+    )
+    assert len(sentence_scores_softmin) == 3
+    assert np.allclose(sentence_scores_softmin, [0.600741787, 1.8005624e-7, 0.8])
+
+    with pytest.raises(AssertionError) as excinfo:
+        get_label_quality_scores(
+            labels, pred_probs, sentence_score_method="unsupported_method", tokens=words
+        )
+    assert "Select from the following methods:" in str(excinfo.value)
+
+
+def test_issues_from_scores(label_quality_scores):
+    sentence_scores, token_scores = label_quality_scores
+    issues = issues_from_scores(sentence_scores, token_scores=token_scores)
+    assert len(issues) == 1
+    assert issues[0] == (1, 0)
+    issues_without = issues_from_scores(sentence_scores)
+    assert len(issues_without) == 1
+    assert issues_without[0] == 1
+
+
+def test_display_issues():
+    display_issues(issues, words)
+    display_issues(issues, tokens=words, labels=labels)
+    display_issues(issues, words, pred_probs=pred_probs)
+    display_issues(issues, words, pred_probs=pred_probs, labels=labels)
+    display_issues(issues, words, pred_probs=pred_probs, labels=labels, class_names=class_names)
+
+    exclude = [(1, 2)]  # Occurs in first token of second sentence "#I"
+    display_issues(issues, words, pred_probs=pred_probs, labels=labels, exclude=exclude)
+
+    top = 1
+    display_issues(issues, words, pred_probs=pred_probs, labels=labels, top=top)
+
+    issues_sentence_only = [i for i, _ in issues]
+    display_issues(issues_sentence_only, words)
+
+
+TEST_KWARGS = {"labels": labels, "pred_probs": pred_probs, "class_names": class_names}
+
+
+@pytest.mark.parametrize(
+    "test_issues",
+    [issues, issues + [(1, 0)]],
+    ids=["default issues", "augmented issues"],
+)
+@pytest.mark.parametrize(
+    "test_kwargs",
+    [
+        {},
+        TEST_KWARGS,
+        {**TEST_KWARGS, "top": 1},
+        {**TEST_KWARGS, "exclude": [(1, 2)]},
+        {**TEST_KWARGS, "verbose": False},
+    ],
+    ids=["no kwargs", "labels+pred_probs+class_names", "...+top", "...+exclude", "...+no verbose"],
+)
+def test_common_label_issues(test_issues, test_kwargs):
+    df = common_label_issues(test_issues, words, **test_kwargs)
+    assert isinstance(df, pd.DataFrame)
+
+    columns = df.columns.tolist()
+    for col in ["token", "num_label_issues"]:
+        assert col in columns
+    if test_kwargs:
+        for col in ["given_label", "predicted_label"]:
+            assert col in columns
+
+
+@pytest.mark.parametrize(
+    "test_token,expected_issues",
+    [
+        ("Hello", []),
+        ("#I", [(1, 0)]),
+    ],
+)
+def test_filter_by_token(test_token, expected_issues):
+    returned_issues = filter_by_token(test_token, issues, words)
+    assert returned_issues == expected_issues

tests/test_util.py

@@ -0,0 +1,213 @@
+# coding: utf-8
+from collections import namedtuple
+
+from cleanlab.internal import util
+import numpy as np
+import pytest
+
+from cleanlab.internal.label_quality_utils import get_normalized_entropy
+from cleanlab.internal.multilabel_utils import int2onehot, onehot2int
+from cleanlab.internal.util import num_unique_classes, format_labels, get_missing_classes
+from cleanlab.internal.validation import assert_valid_class_labels
+
+
+noise_matrix = np.array([[1.0, 0.0, 0.2], [0.0, 0.7, 0.2], [0.0, 0.3, 0.6]])
+
+noise_matrix_2 = np.array(
+    [
+        [1.0, 0.3],
+        [0.0, 0.7],
+    ]
+)
+
+joint_matrix = np.array([[0.1, 0.0, 0.1], [0.1, 0.1, 0.1], [0.2, 0.1, 0.2]])
+
+joint_matrix_2 = np.array(
+    [
+        [0.2, 0.3],
+        [0.4, 0.1],
+    ]
+)
+
+single_element = np.array([1])
+
+
+def test_print_inm():
+    for m in [noise_matrix, noise_matrix_2, single_element]:
+        util.print_inverse_noise_matrix(m, round_places=3)
+
+
+def test_print_joint():
+    for m in [joint_matrix, joint_matrix_2, single_element]:
+        util.print_joint_matrix(m, round_places=3)
+
+
+def test_print_square():
+    for m in [noise_matrix, noise_matrix_2, single_element]:
+        util.print_square_matrix(noise_matrix, round_places=3)
+
+
+def test_print_noise_matrix():
+    for m in [noise_matrix, noise_matrix_2, single_element]:
+        util.print_noise_matrix(noise_matrix, round_places=3)
+
+
+def test_pu_f1():
+    s = [1, 1, 1, 0, 0, 0]
+    p = [1, 1, 1, 0, 0, 0]
+    assert abs(util.estimate_pu_f1(s, p) - 1) < 1e-4
+
+
+def test_value_counts_str():
+    r = util.value_counts(["a", "b", "a"])
+    assert all(np.array([2, 1]) - r < 1e-4)
+
+
+TestCase = namedtuple("TestCase", ["labels", "id"])
+
+value_counts_missing_classes_test_cases = [
+    TestCase([0, 1, 0, 2], "integers"),
+    TestCase(["a", "b", "a", "c"], "strings"),
+    TestCase([[0], [0, 1], [2]], "multilabel_integers"),
+    TestCase([["c"], ["a", "b"], ["a"]], "multilabel_strings"),
+]
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    value_counts_missing_classes_test_cases,
+    ids=lambda x: str(x.id),
+)
+def test_value_counts_fill_missing_classes(test_case):
+    labels = test_case.labels
+    is_multi_label = isinstance(labels[0], list)
+    r = util.value_counts_fill_missing_classes(labels, num_classes=4, multi_label=is_multi_label)
+    assert np.array_equal(r, [2, 1, 1, 0])
+
+
+def test_pu_remove_noise():
+    nm = np.array(
+        [
+            [0.9, 0.0, 0.0],
+            [0.0, 0.7, 0.4],
+            [0.1, 0.3, 0.6],
+        ]
+    )
+    r = util.remove_noise_from_class(nm, 0)
+    assert np.all(r - nm < 1e-4)
+
+
+def test_round_preserving_sum():
+    vec = np.array([1.1] * 10)
+    ints = util.round_preserving_sum(vec)
+    # Make sure one of ints is now 2 to preserve sum of 11
+    assert np.any(ints == 2)
+    assert sum(ints) == 11
+
+
+def test_one_hot():
+    num_classes = 4
+    labels = [[0], [0, 1], [0, 1], [2], [0, 2, 3]]
+    assert onehot2int(int2onehot(labels, K=num_classes)) == labels
+
+
+def test_num_unique():
+    labels = [[0], [0, 1], [0, 1], [2], [0, 2, 3]]
+    assert num_unique_classes(labels) == 4
+
+
+def test_missing_classes():
+    labels = [0, 1]  # class 2 is missing
+    pred_probs = np.array([[0.8, 0.1, 0.1], [0.4, 0.5, 0.1]])
+    assert get_missing_classes(labels, pred_probs=pred_probs) == [2]
+
+
+def test_round_preserving_row_totals():
+    mat = np.array(
+        [
+            [1.7, 1.8, 1.5],
+            [1.1, 1.4, 1.5],
+            [1.3, 1.3, 1.4],
+        ]
+    )
+    mat_int = util.round_preserving_row_totals(mat)
+    # Check that row sums are preserved
+    assert np.all(mat_int.sum(axis=1) == mat.sum(axis=1))
+
+
+def test_confusion_matrix():
+    true = [0, 1, 1, 2, 2, 2]
+    pred = [0, 0, 1, 1, 1, 2]
+    cmat = util.confusion_matrix(true, pred)
+    assert np.shape(cmat) == (3, 3)
+    assert cmat[0][0] == 1
+    assert cmat[1][1] == 1
+    assert cmat[2][2] == 1
+    assert cmat[1][0] == 1
+    assert cmat[2][1] == 2
+    assert cmat[0][1] == 0
+    assert cmat[0][2] == 0
+    assert cmat[2][0] == 0
+    assert cmat[0][1] == 0
+
+
+def test_confusion_matrix_nonconsecutive():
+    true = [-1, -1, -1, 1]
+    pred = [1, 1, -1, 1]
+    cmat = util.confusion_matrix(true, pred)
+    assert np.shape(cmat) == (2, 2)
+    assert cmat[0][0] == 1
+    assert cmat[0][1] == 2
+    assert cmat[1][0] == 0
+    assert cmat[1][1] == 1
+
+
+def test_format_labels():
+    # test 1D labels
+    str_labels = np.array(["b", "b", "a", "c", "a"])
+    labels, label_map = format_labels(str_labels)
+    assert all(labels == np.array([1, 1, 0, 2, 0]))
+    assert label_map[0] == "a"
+    assert label_map[1] == "b"
+    assert label_map[2] == "c"
+    assert_valid_class_labels(labels)
+
+
+def test_normalized_entropy():
+    """Check that normalized entropy is well well-behaved and in [0, 1]."""
+    # test tiny numbers
+    for dtype in [np.float16, np.float32, np.float64]:
+        info = np.finfo(dtype)
+        # some NumPy versions have bugs, therefore we provide a fallback
+        # (fallback is the value of the smalles datatype float16)
+        smallest_normal = getattr(info, "smallest_normal", 6.104e-05)
+        smallest_subnormal = getattr(info, "smallest_subnormal", 6e-08)
+        for val in [info.eps, smallest_normal, smallest_subnormal, 0]:
+            entropy = get_normalized_entropy(np.array([[1.0, val]], dtype=dtype))
+            assert 0.0 <= entropy <= 1.0
+    # test multiple _assert_valid_inputs
+    entropy = get_normalized_entropy(np.array([[0.0, 1.0], [0.5, 0.5]]))
+    assert all((0.0 <= entropy) & (entropy <= 1.0))
+
+    # raise errors for wrong probabilities.
+    with pytest.raises(ValueError):
+        get_normalized_entropy(np.array([[-1.0, 0.5]]))  # negative
+        get_normalized_entropy(np.array([[2.0, 0.5]]))  # larger 1
+
+
+def test_force_two_dimensions():
+    # Test with 2D array
+    X = np.zeros((5, 5))
+    X_reshaped = util.force_two_dimensions(X)
+    assert X_reshaped.shape == (5, 5), "The shape of 2D array should remain unchanged."
+
+    # Test with 4D array
+    X = np.zeros((5, 5, 5, 5))
+    X_reshaped = util.force_two_dimensions(X)
+    assert X_reshaped.shape == (
+        5,
+        125,
+    ), "The shape of 4D array should be flattened to two dimensions."
+
+    # Test with None input
+    assert util.force_two_dimensions(None) is None, "None input should return None."

tests/test_validation.py

@@ -0,0 +1,29 @@
+# coding: utf-8
+
+from cleanlab.internal import validation
+import numpy as np
+import pandas as pd
+import pytest
+
+
+@pytest.mark.parametrize("y_list", [["a", "b", "a"], [0, 1, 2]])
+@pytest.mark.parametrize("format", [list, np.array, pd.Series, pd.DataFrame])
+def test_labels_to_array_return_types(y_list, format):
+    y = format(y_list)
+    labels = validation.labels_to_array(y)
+    assert isinstance(labels, np.ndarray)
+
+
+@pytest.mark.parametrize("y_list", [["a", "b", "a"], [0, 1, 2]])
+@pytest.mark.parametrize("format", [list, np.array, pd.Series])
+def test_labels_to_array_return_values(y_list, format):
+    y = format(y_list)
+    labels = validation.labels_to_array(y)
+    assert np.array_equal(y, labels)
+
+
+def test_label_to_array_raises_error():
+    # Pandas DataFrame should have only one column
+    y = pd.DataFrame({"a": [0, 1], "b": [2, 3]})
+    with pytest.raises(ValueError):
+        validation.labels_to_array(y)