import numpy as np
from sklearn.metrics import multilabel_confusion_matrix as skmultilabel_confusion_matrix

def multilabel_confusion_matrix(y_true, y_pred):
    n_labels = len(set(y_true) | set(y_pred))
    true_sum = np.bincount(y_true, minlength=n_labels)
    pred_sum = np.bincount(y_pred, minlength=n_labels)
    tp = np.bincount(y_true[y_true == y_pred], minlength=n_labels)
    fp = pred_sum - tp
    fn = true_sum - tp
    tn = len(y_true) - tp - fp - fn
    return np.array([tn, fp, fn, tp]).T.reshape(-1, 2, 2)

# binary
for i in range(10):
    rng = np.random.RandomState(i)
    y_true = rng.randint(2, size=10)
    y_pred = rng.randint(2, size=10)
    score1 = multilabel_confusion_matrix(y_true, y_pred)
    score2 = skmultilabel_confusion_matrix(y_true, y_pred)
    assert np.array_equal(score1, score2)

# multiclass
for i in range(10):
    rng = np.random.RandomState(i)
    y_true = rng.randint(3, size=10)
    y_pred = rng.randint(3, size=10)
    score1 = multilabel_confusion_matrix(y_true, y_pred)
    score2 = skmultilabel_confusion_matrix(y_true, y_pred)
    assert np.array_equal(score1, score2)