type of auc curve roc or pr roc roc pr pr staticmethod def fromstrkey if key in pr pr return auccurve pr elif key in roc roc return auccurve roc else raise valueerror f invalid auc curve value key expected values are pr roc class aucsummationmethodenum interpolation interpolation majoring majoring minoring minoring staticmethod def fromstrkey if key in interpolation interpolation return aucsummationmethod interpolation elif key in majoring majoring return aucsummationmethod majoring elif key in minoring minoring return aucsummationmethod minoring else raise valueerror f invalid auc summation method value key expected values are interpolation majoring minoring def updateconfusionmatrixvariablesoptimized variablestoupdate ytrue ypred thresholds multilabelfalse sampleweightsnone labelweightsnone thresholdswithepsilonfalse the meaning of this bucket is that if any of the label is true then 1 will be added to the corresponding bucket with the index eg if the label for 0 2 is true then 1 will be added to bucket 0 if the label for 1 8 is true then 1 will be added to bucket 1 note the second item 1 0 is floored to 0 since the value need to be strictly larger than the bucket lower bound in the implementation we use ops ceil 1 to achieve this for 1 1 0 here it means there is 1 true value contributed by bucket 0 and 1 value contributed by bucket 1 when we aggregate them to together the result become a b c b c c since large thresholds will always contribute to the value for smaller thresholds numthresholds ops shapethresholds0 if sampleweights is none sampleweights 1 0 else sampleweights ops broadcastto ops castsampleweights dtypeypred dtype ops shapeypred if not multilabel sampleweights ops reshapesampleweights 1 if labelweights is none labelweights 1 0 else labelweights ops expanddimslabelweights 0 labelweights ops broadcasttolabelweights ops shapeypred if not multilabel labelweights ops reshapelabelweights 1 weights ops cast ops multiplysampleweights labelweights ytrue dtype we shouldn t need this but in case there are predict value that is out of the range of 0 0 1 0 ypred ops clipypred xmin0 0 xmax1 0 ytrue ops castops castytrue bool ytrue dtype if not multilabel ytrue ops reshapeytrue 1 ypred ops reshapeypred 1 truelabels ops multiplyytrue weights falselabels ops multiply1 0 ytrue weights compute the bucket indices for each prediction value since the predict value has to be strictly greater than the thresholds eg buckets like 0 0 5 0 5 1 and 0 5 belongs to first bucket we have to use math ceilval 1 for the bucket bucketindices ops ceilypred ops castnumthresholds dtypeypred dtype 1 1 if thresholdswithepsilon in this case the first bucket should actually take into account since the any prediction between 0 0 1 0 should be larger than the first threshold we change the bucket value from 1 to 0 bucketindices ops relubucketindices bucketindices ops castbucketindices int32 if multilabel we need to run bucket segment sum for each of the label class in the multilabel case the rank of the label is 2 we first transpose it so that the label dim becomes the first and we can parallel run though them truelabels ops transposetruelabels falselabels ops transposefalselabels bucketindices ops transposebucketindices def gatherbucketlabelandbucketindex label bucketindex labelandbucketindex0 labelandbucketindex1 return ops segmentsum datalabel segmentidsbucketindex numsegmentsnumthresholds tpbucketv backend vectorizedmap gatherbucket truelabels bucketindices fpbucketv backend vectorizedmap gatherbucket falselabels bucketindices tp ops transposeops flipops cumsumops fliptpbucketv axis1 fp ops transposeops flipops cumsumops flipfpbucketv axis1 else tpbucketv ops segmentsum datatruelabels segmentidsbucketindices numsegmentsnumthresholds fpbucketv ops segmentsum datafalselabels segmentidsbucketindices numsegmentsnumthresholds tp ops flipops cumsumops fliptpbucketv fp ops flipops cumsumops flipfpbucketv fn sumtruelabels tp tn sumfalselabels fp if confusionmatrix truenegatives in variablestoupdate or confusionmatrix falsenegatives in variablestoupdate if multilabel totaltruelabels ops sumtruelabels axis1 totalfalselabels ops sumfalselabels axis1 else totaltruelabels ops sumtruelabels totalfalselabels ops sumfalselabels if confusionmatrix truepositives in variablestoupdate variable variablestoupdateconfusionmatrix truepositives variable assignvariable tp if confusionmatrix falsepositives in variablestoupdate variable variablestoupdateconfusionmatrix falsepositives variable assignvariable fp if confusionmatrix truenegatives in variablestoupdate variable variablestoupdateconfusionmatrix truenegatives tn totalfalselabels fp variable assignvariable tn if confusionmatrix falsenegatives in variablestoupdate variable variablestoupdateconfusionmatrix falsenegatives fn totaltruelabels tp variable assignvariable fn def isevenlydistributedthresholdsthresholds check the list value and see if it is evenly distributed numthresholds lenthresholds if numthresholds 3 return false eventhresholds np arangenumthresholds dtypenp float32 numthresholds 1 return np allclosethresholds eventhresholds atolbackend epsilon def updateconfusionmatrixvariables variablestoupdate ytrue ypred thresholds topknone classidnone sampleweightnone multilabelfalse labelweightsnone thresholdsdistributedevenlyfalse if multilabel and labelweights is not none raise valueerror labelweights for multilabel data should be handled outside of updateconfusionmatrixvariables when multilabel is true if variablestoupdate is none return if not any key for key in variablestoupdate if key in listconfusionmatrix raise valueerror please provide at least one valid confusion matrix variable to update valid variable key options are f listconfusionmatrix f received variablestoupdate keys variabledtype listvariablestoupdate values0 dtype ytrue ops castytrue dtypevariabledtype ypred ops castypred dtypevariabledtype if thresholdsdistributedevenly check whether the thresholds has any leading or tailing epsilon added for floating point imprecision the leading and tailing threshold will be handled bit differently as the corner case at this point thresholds should be a listarray with more than 2 items and ranged between 0 1 see isevenlydistributedthresholds for more details thresholdswithepsilon thresholds0 0 0 or thresholds1 1 0 thresholds ops converttotensorthresholds dtypevariabledtype numthresholds ops shapethresholds0 if multilabel onethresh ops equal np array1 dtypeint32 lenthresholds shape else onethresh np arraytrue dtypebool invalidkeys key for key in variablestoupdate if key not in listconfusionmatrix if invalidkeys raise valueerror f invalid keys invalidkeys f valid variable key options are listconfusionmatrix ypred ytrue squeezetosamerankypred ytrue if sampleweight is not none sampleweight ops expanddims ops castsampleweight dtypevariabledtype axis1 sampleweight squeezetosamerankytrue sampleweight if topk is not none ypred filtertopkypred topk if classid is not none if lenypred shape 1 raise valueerror when classid is provided ypred must be a 2d array with shape numsamples numclasses found shape fypred shape preserve dimension to match with sampleweight ytrue ytrue classid none ypred ypred classid none if thresholdsdistributedevenly return updateconfusionmatrixvariablesoptimized variablestoupdate ytrue ypred thresholds multilabelmultilabel sampleweightssampleweight labelweightslabelweights thresholdswithepsilonthresholdswithepsilon if none in ypred shape predshape ops shapeypred numpredictions predshape0 if lenypred shape 1 numlabels 1 else numlabels ops cast ops prodops arraypredshape1 axis0 int32 threshlabeltile ops whereonethresh numlabels 1 else predshape ypred shape numpredictions predshape0 if lenypred shape 1 numlabels 1 else numlabels np prodpredshape1 axis0 astypeint32 threshlabeltile np whereonethresh numlabels 1 reshape predictions and labels adding a dim for thresholding if multilabel predictionsextradim ops expanddimsypred 0 labelsextradim ops expanddimsops castytrue dtypebool 0 else flatten predictions and labels when not multilabel predictionsextradim ops reshapeypred 1 1 labelsextradim ops reshapeops castytrue dtypebool 1 1 tile the thresholds for every prediction if multilabel threshpretileshape numthresholds 1 1 threshtiles 1 numpredictions threshlabeltile datatiles numthresholds 1 1 else threshpretileshape numthresholds 1 threshtiles 1 numpredictions numlabels datatiles numthresholds 1 threshtiled ops tile ops reshapethresholds threshpretileshape threshtiles tile the predictions for every threshold predstiled ops tilepredictionsextradim datatiles compare predictions and threshold predispos ops greaterpredstiled threshtiled tile labels by number of thresholds labelispos ops tilelabelsextradim datatiles if sampleweight is not none sampleweight ops broadcastto ops castsampleweight dtypeypred dtype ops shapeypred weightstiled ops tile ops reshapesampleweight threshtiles datatiles else weightstiled none if labelweights is not none and not multilabel labelweights ops expanddimslabelweights 0 labelweights ops broadcasttolabelweights ops shapeypred labelweightstiled ops tile ops reshapelabelweights threshtiles datatiles if weightstiled is none weightstiled labelweightstiled else weightstiled ops multiplyweightstiled labelweightstiled def weightedassignaddlabel pred weights var labelandpred ops castops logicalandlabel pred dtypevar dtype if weights is not none labelandpred ops castweights dtypevar dtype var assignvar ops sumlabelandpred 1 loopvars confusionmatrix truepositives labelispos predispos updatetn confusionmatrix truenegatives in variablestoupdate updatefp confusionmatrix falsepositives in variablestoupdate updatefn confusionmatrix falsenegatives in variablestoupdate if updatefn or updatetn predisneg ops logicalnotpredispos loopvarsconfusionmatrix falsenegatives labelispos predisneg if updatefp or updatetn labelisneg ops logicalnotlabelispos loopvarsconfusionmatrix falsepositives labelisneg predispos if updatetn loopvarsconfusionmatrix truenegatives labelisneg predisneg for matrixcond label pred in loopvars items if matrixcond in variablestoupdate weightedassignadd label pred weightstiled variablestoupdatematrixcond def filtertopkx k topkidx ops topkx k topkmask ops sum ops onehottopkidx ops shapex1 axis1 axis2 return x topkmask neginf 1 topkmask def confusionmatrix labels predictions numclasses weightsnone dtypeint32 labels ops converttotensorlabels dtype predictions ops converttotensorpredictions dtype labels predictions squeezetosameranklabels predictions predictions ops castpredictions dtype labels ops castlabels dtype if weights is not none weights ops converttotensorweights dtype indices ops stacklabels predictions axis1 values ops oneslikepredictions dtype if weights is none else weights indices ops castindices dtypeint64 values ops castvalues dtypedtype numclasses intnumclasses confusionmatrix ops scatterindices values numclasses numclasses return confusionmatrix type of auc curve roc or pr type of auc summation method https en wikipedia org wiki riemann_sum contains the following values interpolation applies mid point summation scheme for roc curve for pr curve interpolates true false positives but not the ratio that is precision see davis goadrich 2006 for details minoring applies left summation for increasing intervals and right summation for decreasing intervals majoring applies right summation for increasing intervals and left summation for decreasing intervals update confusion matrix variables with memory efficient alternative note that the thresholds need to be evenly distributed within the list eg the diff between consecutive elements are the same to compute tp fp tn fn we are measuring a binary classifier c t predictions t at each threshold t so we have tp t sum c t true_labels fp t sum c t false_labels but computing c t requires computation for each t to make it fast observe that c t is a cumulative integral and so if we have thresholds t_0 t_ n 1 t_0 t_ n 1 where n num_thresholds and if we can compute the bucket function b i sum predictions t t_i t t i 1 then we get c t_i sum b j j i which is the reversed cumulative sum in ops cumsum we can compute b i efficiently by taking advantage of the fact that our thresholds are evenly distributed in that width 1 0 num_thresholds 1 thresholds 0 0 1 width 2 width 3 width 1 0 given a prediction value p we can map it to its bucket by bucket_index p floor p num_thresholds 1 so we can use ops segment_sum to update the buckets in one pass consider following example y_true 0 0 1 1 y_pred 0 1 0 5 0 3 0 9 thresholds 0 0 0 5 1 0 num_buckets 2 0 0 1 0 1 0 2 0 bucket_index y_pred ops floor y_pred num_buckets ops floor 0 2 1 0 0 6 1 8 0 0 0 1 the meaning of this bucket is that if any of the label is true then 1 will be added to the corresponding bucket with the index eg if the label for 0 2 is true then 1 will be added to bucket 0 if the label for 1 8 is true then 1 will be added to bucket 1 note the second item 1 0 is floored to 0 since the value need to be strictly larger than the bucket lower bound in the implementation we use ops ceil 1 to achieve this tp_bucket_value ops segment_sum true_labels bucket_indices num_segments num_thresholds 1 1 0 for 1 1 0 here it means there is 1 true value contributed by bucket 0 and 1 value contributed by bucket 1 when we aggregate them to together the result become a b c b c c since large thresholds will always contribute to the value for smaller thresholds true_positive ops cumsum tp_bucket_value reverse true 2 1 0 this implementation exhibits a run time and space complexity of o t n where t is the number of thresholds and n is the size of predictions metrics that rely on standard implementation instead exhibit a complexity of o t n args variables_to_update dictionary with tp fn tn fp as valid keys and corresponding variables to update as values y_true a floating point tensor whose shape matches y_pred will be cast to bool y_pred a floating point tensor of arbitrary shape and whose values are in the range 0 1 thresholds a sorted floating point tensor with value in 0 1 it need to be evenly distributed the diff between each element need to be the same multi_label optional boolean indicating whether multidimensional prediction labels should be treated as multilabel responses or flattened into a single label when true the valus of variables_to_update must have a second dimension equal to the number of labels in y_true and y_pred and those tensors must not be raggedtensors sample_weights optional tensor whose rank is either 0 or the same rank as y_true and must be broadcastable to y_true i e all dimensions must be either 1 or the same as the corresponding y_true dimension label_weights optional tensor of non negative weights for multilabel data the weights are applied when calculating tp fp fn and tn without explicit multilabel handling i e when the data is to be flattened thresholds_with_epsilon optional boolean indicating whether the leading and tailing thresholds has any epsilon added for floating point imprecisions it will change how we handle the leading and tailing bucket we shouldn t need this but in case there are predict value that is out of the range of 0 0 1 0 compute the bucket indices for each prediction value since the predict value has to be strictly greater than the thresholds eg buckets like 0 0 5 0 5 1 and 0 5 belongs to first bucket we have to use math ceil val 1 for the bucket in this case the first bucket should actually take into account since the any prediction between 0 0 1 0 should be larger than the first threshold we change the bucket value from 1 to 0 we need to run bucket segment sum for each of the label class in the multi_label case the rank of the label is 2 we first transpose it so that the label dim becomes the first and we can parallel run though them fn sum true_labels tp tn sum false_labels fp check if the thresholds list is evenly distributed we could leverage evenly distributed thresholds to use less memory when calculate metrcis like auc where each individual threshold need to be evaluated args thresholds a python list or tuple or 1d numpy array whose value is ranged in 0 1 returns boolean whether the values in the inputs are evenly distributed check the list value and see if it is evenly distributed updates the given confusion matrix variables for every pair of values in y_true and y_pred true_positive y_true true and y_pred thresholds false_negatives y_true true and y_pred thresholds true_negatives y_true false and y_pred thresholds false_positive y_true false and y_pred thresholds the results will be weighted and added together when multiple thresholds are provided we will repeat the same for every threshold for estimation of these metrics over a stream of data the function creates an update_op operation that updates the given variables if sample_weight is none weights default to 1 use weights of 0 to mask values args variables_to_update dictionary with tp fn tn fp as valid keys and corresponding variables to update as values y_true a tensor whose shape matches y_pred will be cast to bool y_pred a floating point tensor of arbitrary shape and whose values are in the range 0 1 thresholds a float value float tensor python list or tuple of float thresholds in 0 1 or neg_inf used when top_k is set top_k optional int indicates that the positive labels should be limited to the top k predictions class_id optional int limits the prediction and labels to the class specified by this argument sample_weight optional tensor whose rank is either 0 or the same rank as y_true and must be broadcastable to y_true i e all dimensions must be either 1 or the same as the corresponding y_true dimension multi_label optional boolean indicating whether multidimensional prediction labels should be treated as multilabel responses or flattened into a single label when true the valus of variables_to_update must have a second dimension equal to the number of labels in y_true and y_pred and those tensors must not be raggedtensors label_weights optional tensor of non negative weights for multilabel data the weights are applied when calculating tp fp fn and tn without explicit multilabel handling i e when the data is to be flattened thresholds_distributed_evenly boolean whether the thresholds are evenly distributed within the list an optimized method will be used if this is the case see _update_confusion_matrix_variables_optimized for more details raises valueerror if y_pred and y_true have mismatched shapes or if sample_weight is not none and its shape doesn t match y_pred or if variables_to_update contains invalid keys check whether the thresholds has any leading or tailing epsilon added for floating point imprecision the leading and tailing threshold will be handled bit differently as the corner case at this point thresholds should be a list array with more than 2 items and ranged between 0 1 see is_evenly_distributed_thresholds for more details preserve dimension to match with sample_weight reshape predictions and labels adding a dim for thresholding flatten predictions and labels when not multilabel tile the thresholds for every prediction tile the predictions for every threshold compare predictions and threshold tile labels by number of thresholds filters top k values in the last dim of x and set the rest to neg_inf used for computing top k prediction values in dense labels which has the same shape as predictions for recall and precision top k metrics args x tensor with any dimensions k the number of values to keep returns tensor with same shape and dtype as x computes the confusion matrix from predictions and labels the matrix columns represent the prediction labels and the rows represent the real labels the confusion matrix is always a 2 d array of shape n n where n is the number of valid labels for a given classification task both prediction and labels must be 1 d arrays of the same shape in order for this function to work if num_classes is none then num_classes will be set to one plus the maximum value in either predictions or labels class labels are expected to start at 0 for example if num_classes is 3 then the possible labels would be 0 1 2 if weights is not none then each prediction contributes its corresponding weight to the total value of the confusion matrix cell for example python keras metrics metrics_utils confusion_matrix 1 2 4 2 2 4 0 0 0 0 0 0 0 1 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 1 note that the possible labels are assumed to be 0 1 2 3 4 resulting in a 5x5 confusion matrix args labels 1 d tensor of real labels for the classification task predictions 1 d tensor of predictions for a given classification num_classes the possible number of labels the classification task can have weights an optional tensor whose shape matches predictions dtype data type of the confusion matrix returns a tensor of type dtype with shape n n representing the confusion matrix where n is the number of possible labels in the classification task
from enum import Enum

import numpy as np

from keras import backend
from keras import ops
from keras.losses.loss import squeeze_to_same_rank
from keras.utils.python_utils import to_list

NEG_INF = -1e10


def assert_thresholds_range(thresholds):
    if thresholds is not None:
        invalid_thresholds = [
            t for t in thresholds if t is None or t < 0 or t > 1
        ]
        if invalid_thresholds:
            raise ValueError(
                "Threshold values must be in [0, 1]. "
                f"Received: {invalid_thresholds}"
            )


def parse_init_thresholds(thresholds, default_threshold=0.5):
    if thresholds is not None:
        assert_thresholds_range(to_list(thresholds))
    thresholds = to_list(
        default_threshold if thresholds is None else thresholds
    )
    return thresholds


class ConfusionMatrix(Enum):
    TRUE_POSITIVES = "tp"
    FALSE_POSITIVES = "fp"
    TRUE_NEGATIVES = "tn"
    FALSE_NEGATIVES = "fn"


class AUCCurve(Enum):
    

    ROC = "ROC"
    PR = "PR"

    @staticmethod
    def from_str(key):
        if key in ("pr", "PR"):
            return AUCCurve.PR
        elif key in ("roc", "ROC"):
            return AUCCurve.ROC
        else:
            raise ValueError(
                f'Invalid AUC curve value: "{key}". '
                'Expected values are ["PR", "ROC"]'
            )


class AUCSummationMethod(Enum):
    

    INTERPOLATION = "interpolation"
    MAJORING = "majoring"
    MINORING = "minoring"

    @staticmethod
    def from_str(key):
        if key in ("interpolation", "Interpolation"):
            return AUCSummationMethod.INTERPOLATION
        elif key in ("majoring", "Majoring"):
            return AUCSummationMethod.MAJORING
        elif key in ("minoring", "Minoring"):
            return AUCSummationMethod.MINORING
        else:
            raise ValueError(
                f'Invalid AUC summation method value: "{key}". '
                'Expected values are ["interpolation", "majoring", "minoring"]'
            )


def _update_confusion_matrix_variables_optimized(
    variables_to_update,
    y_true,
    y_pred,
    thresholds,
    multi_label=False,
    sample_weights=None,
    label_weights=None,
    thresholds_with_epsilon=False,
):
    
    num_thresholds = ops.shape(thresholds)[0]

    if sample_weights is None:
        sample_weights = 1.0
    else:
        sample_weights = ops.broadcast_to(
            ops.cast(sample_weights, dtype=y_pred.dtype), ops.shape(y_pred)
        )
        if not multi_label:
            sample_weights = ops.reshape(sample_weights, [-1])
    if label_weights is None:
        label_weights = 1.0
    else:
        label_weights = ops.expand_dims(label_weights, 0)
        label_weights = ops.broadcast_to(label_weights, ops.shape(y_pred))
        if not multi_label:
            label_weights = ops.reshape(label_weights, [-1])
    weights = ops.cast(
        ops.multiply(sample_weights, label_weights), y_true.dtype
    )

    
    
    y_pred = ops.clip(y_pred, x_min=0.0, x_max=1.0)

    y_true = ops.cast(ops.cast(y_true, "bool"), y_true.dtype)
    if not multi_label:
        y_true = ops.reshape(y_true, [-1])
        y_pred = ops.reshape(y_pred, [-1])

    true_labels = ops.multiply(y_true, weights)
    false_labels = ops.multiply((1.0 - y_true), weights)

    
    
    
    
    bucket_indices = (
        ops.ceil(y_pred * (ops.cast(num_thresholds, dtype=y_pred.dtype) - 1))
        - 1
    )

    if thresholds_with_epsilon:
        
        
        
        bucket_indices = ops.relu(bucket_indices)

    bucket_indices = ops.cast(bucket_indices, "int32")

    if multi_label:
        
        
        
        
        true_labels = ops.transpose(true_labels)
        false_labels = ops.transpose(false_labels)
        bucket_indices = ops.transpose(bucket_indices)

        def gather_bucket(label_and_bucket_index):
            label, bucket_index = (
                label_and_bucket_index[0],
                label_and_bucket_index[1],
            )
            return ops.segment_sum(
                data=label,
                segment_ids=bucket_index,
                num_segments=num_thresholds,
            )

        tp_bucket_v = backend.vectorized_map(
            gather_bucket,
            (true_labels, bucket_indices),
        )
        fp_bucket_v = backend.vectorized_map(
            gather_bucket, (false_labels, bucket_indices)
        )
        tp = ops.transpose(ops.flip(ops.cumsum(ops.flip(tp_bucket_v), axis=1)))
        fp = ops.transpose(ops.flip(ops.cumsum(ops.flip(fp_bucket_v), axis=1)))
    else:
        tp_bucket_v = ops.segment_sum(
            data=true_labels,
            segment_ids=bucket_indices,
            num_segments=num_thresholds,
        )
        fp_bucket_v = ops.segment_sum(
            data=false_labels,
            segment_ids=bucket_indices,
            num_segments=num_thresholds,
        )
        tp = ops.flip(ops.cumsum(ops.flip(tp_bucket_v)))
        fp = ops.flip(ops.cumsum(ops.flip(fp_bucket_v)))

    
    
    if (
        ConfusionMatrix.TRUE_NEGATIVES in variables_to_update
        or ConfusionMatrix.FALSE_NEGATIVES in variables_to_update
    ):
        if multi_label:
            total_true_labels = ops.sum(true_labels, axis=1)
            total_false_labels = ops.sum(false_labels, axis=1)
        else:
            total_true_labels = ops.sum(true_labels)
            total_false_labels = ops.sum(false_labels)

    if ConfusionMatrix.TRUE_POSITIVES in variables_to_update:
        variable = variables_to_update[ConfusionMatrix.TRUE_POSITIVES]
        variable.assign(variable + tp)
    if ConfusionMatrix.FALSE_POSITIVES in variables_to_update:
        variable = variables_to_update[ConfusionMatrix.FALSE_POSITIVES]
        variable.assign(variable + fp)
    if ConfusionMatrix.TRUE_NEGATIVES in variables_to_update:
        variable = variables_to_update[ConfusionMatrix.TRUE_NEGATIVES]
        tn = total_false_labels - fp
        variable.assign(variable + tn)
    if ConfusionMatrix.FALSE_NEGATIVES in variables_to_update:
        variable = variables_to_update[ConfusionMatrix.FALSE_NEGATIVES]
        fn = total_true_labels - tp
        variable.assign(variable + fn)


def is_evenly_distributed_thresholds(thresholds):
    
    
    num_thresholds = len(thresholds)
    if num_thresholds < 3:
        return False
    even_thresholds = np.arange(num_thresholds, dtype=np.float32) / (
        num_thresholds - 1
    )
    return np.allclose(thresholds, even_thresholds, atol=backend.epsilon())


def update_confusion_matrix_variables(
    variables_to_update,
    y_true,
    y_pred,
    thresholds,
    top_k=None,
    class_id=None,
    sample_weight=None,
    multi_label=False,
    label_weights=None,
    thresholds_distributed_evenly=False,
):
    
    if multi_label and label_weights is not None:
        raise ValueError(
            "`label_weights` for multilabel data should be handled "
            "outside of `update_confusion_matrix_variables` when "
            "`multi_label` is True."
        )
    if variables_to_update is None:
        return
    if not any(
        key for key in variables_to_update if key in list(ConfusionMatrix)
    ):
        raise ValueError(
            "Please provide at least one valid confusion matrix "
            "variable to update. Valid variable key options are: "
            f'"{list(ConfusionMatrix)}". '
            f'Received: "{variables_to_update.keys()}"'
        )

    variable_dtype = list(variables_to_update.values())[0].dtype

    y_true = ops.cast(y_true, dtype=variable_dtype)
    y_pred = ops.cast(y_pred, dtype=variable_dtype)

    if thresholds_distributed_evenly:
        
        
        
        
        
        
        thresholds_with_epsilon = thresholds[0] < 0.0 or thresholds[-1] > 1.0

    thresholds = ops.convert_to_tensor(thresholds, dtype=variable_dtype)
    num_thresholds = ops.shape(thresholds)[0]

    if multi_label:
        one_thresh = ops.equal(
            np.array(1, dtype="int32"),
            len(thresholds.shape),
        )
    else:
        one_thresh = np.array(True, dtype="bool")

    invalid_keys = [
        key for key in variables_to_update if key not in list(ConfusionMatrix)
    ]
    if invalid_keys:
        raise ValueError(
            f'Invalid keys: "{invalid_keys}". '
            f'Valid variable key options are: "{list(ConfusionMatrix)}"'
        )

    y_pred, y_true = squeeze_to_same_rank(y_pred, y_true)
    if sample_weight is not None:
        sample_weight = ops.expand_dims(
            ops.cast(sample_weight, dtype=variable_dtype), axis=-1
        )
        _, sample_weight = squeeze_to_same_rank(y_true, sample_weight)

    if top_k is not None:
        y_pred = _filter_top_k(y_pred, top_k)

    if class_id is not None:
        if len(y_pred.shape) == 1:
            raise ValueError(
                "When class_id is provided, y_pred must be a 2D array "
                "with shape (num_samples, num_classes), found shape: "
                f"{y_pred.shape}"
            )

        
        y_true = y_true[..., class_id, None]
        y_pred = y_pred[..., class_id, None]

    if thresholds_distributed_evenly:
        return _update_confusion_matrix_variables_optimized(
            variables_to_update,
            y_true,
            y_pred,
            thresholds,
            multi_label=multi_label,
            sample_weights=sample_weight,
            label_weights=label_weights,
            thresholds_with_epsilon=thresholds_with_epsilon,
        )

    if None in y_pred.shape:
        pred_shape = ops.shape(y_pred)
        num_predictions = pred_shape[0]
        if len(y_pred.shape) == 1:
            num_labels = 1
        else:
            num_labels = ops.cast(
                ops.prod(ops.array(pred_shape[1:]), axis=0), "int32"
            )
        thresh_label_tile = ops.where(one_thresh, num_labels, 1)
    else:
        pred_shape = y_pred.shape
        num_predictions = pred_shape[0]
        if len(y_pred.shape) == 1:
            num_labels = 1
        else:
            num_labels = np.prod(pred_shape[1:], axis=0).astype("int32")
        thresh_label_tile = np.where(one_thresh, num_labels, 1)

    
    if multi_label:
        predictions_extra_dim = ops.expand_dims(y_pred, 0)
        labels_extra_dim = ops.expand_dims(ops.cast(y_true, dtype="bool"), 0)
    else:
        
        predictions_extra_dim = ops.reshape(y_pred, [1, -1])
        labels_extra_dim = ops.reshape(ops.cast(y_true, dtype="bool"), [1, -1])

    
    if multi_label:
        thresh_pretile_shape = [num_thresholds, 1, -1]
        thresh_tiles = [1, num_predictions, thresh_label_tile]
        data_tiles = [num_thresholds, 1, 1]
    else:
        thresh_pretile_shape = [num_thresholds, -1]
        thresh_tiles = [1, num_predictions * num_labels]
        data_tiles = [num_thresholds, 1]

    thresh_tiled = ops.tile(
        ops.reshape(thresholds, thresh_pretile_shape), thresh_tiles
    )

    
    preds_tiled = ops.tile(predictions_extra_dim, data_tiles)

    
    pred_is_pos = ops.greater(preds_tiled, thresh_tiled)

    
    label_is_pos = ops.tile(labels_extra_dim, data_tiles)

    if sample_weight is not None:
        sample_weight = ops.broadcast_to(
            ops.cast(sample_weight, dtype=y_pred.dtype), ops.shape(y_pred)
        )
        weights_tiled = ops.tile(
            ops.reshape(sample_weight, thresh_tiles), data_tiles
        )
    else:
        weights_tiled = None

    if label_weights is not None and not multi_label:
        label_weights = ops.expand_dims(label_weights, 0)
        label_weights = ops.broadcast_to(label_weights, ops.shape(y_pred))
        label_weights_tiled = ops.tile(
            ops.reshape(label_weights, thresh_tiles), data_tiles
        )
        if weights_tiled is None:
            weights_tiled = label_weights_tiled
        else:
            weights_tiled = ops.multiply(weights_tiled, label_weights_tiled)

    def weighted_assign_add(label, pred, weights, var):
        label_and_pred = ops.cast(ops.logical_and(label, pred), dtype=var.dtype)
        if weights is not None:
            label_and_pred *= ops.cast(weights, dtype=var.dtype)
        var.assign(var + ops.sum(label_and_pred, 1))

    loop_vars = {
        ConfusionMatrix.TRUE_POSITIVES: (label_is_pos, pred_is_pos),
    }
    update_tn = ConfusionMatrix.TRUE_NEGATIVES in variables_to_update
    update_fp = ConfusionMatrix.FALSE_POSITIVES in variables_to_update
    update_fn = ConfusionMatrix.FALSE_NEGATIVES in variables_to_update

    if update_fn or update_tn:
        pred_is_neg = ops.logical_not(pred_is_pos)
        loop_vars[ConfusionMatrix.FALSE_NEGATIVES] = (label_is_pos, pred_is_neg)

    if update_fp or update_tn:
        label_is_neg = ops.logical_not(label_is_pos)
        loop_vars[ConfusionMatrix.FALSE_POSITIVES] = (label_is_neg, pred_is_pos)
        if update_tn:
            loop_vars[ConfusionMatrix.TRUE_NEGATIVES] = (
                label_is_neg,
                pred_is_neg,
            )

    for matrix_cond, (label, pred) in loop_vars.items():
        if matrix_cond in variables_to_update:
            weighted_assign_add(
                label, pred, weights_tiled, variables_to_update[matrix_cond]
            )


def _filter_top_k(x, k):
    
    _, top_k_idx = ops.top_k(x, k)
    top_k_mask = ops.sum(
        ops.one_hot(top_k_idx, ops.shape(x)[-1], axis=-1), axis=-2
    )
    return x * top_k_mask + NEG_INF * (1 - top_k_mask)


def confusion_matrix(
    labels,
    predictions,
    num_classes,
    weights=None,
    dtype="int32",
):
    
    labels = ops.convert_to_tensor(labels, dtype)
    predictions = ops.convert_to_tensor(predictions, dtype)
    labels, predictions = squeeze_to_same_rank(labels, predictions)

    predictions = ops.cast(predictions, dtype)
    labels = ops.cast(labels, dtype)

    if weights is not None:
        weights = ops.convert_to_tensor(weights, dtype)

    indices = ops.stack([labels, predictions], axis=1)
    values = ops.ones_like(predictions, dtype) if weights is None else weights
    indices = ops.cast(indices, dtype="int64")
    values = ops.cast(values, dtype=dtype)
    num_classes = int(num_classes)
    confusion_matrix = ops.scatter(indices, values, (num_classes, num_classes))
    return confusion_matrix