code: 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,
  ):
      """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.
      """
      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
      )
  
      # We shouldn't need this, but in case there are predict value that is out of
      # the range of [0.0, 1.0]
      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)
  
      # 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.
      bucket_indices = (
          ops.ceil(
              y_pred
              * (ops.cast(num_thresholds, dtype=y_pred.dtype) - 1)
          )
          - 1
      )
  
      if thresholds_with_epsilon:
          # 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.
          bucket_indices = ops.relu(bucket_indices)
  
      bucket_indices = ops.cast(bucket_indices, "int32")
  
      if multi_label:
          # 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.
          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)))
  
      # fn = sum(true_labels) - tp
      # tn = sum(false_labels) - fp
      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)
  
imports: from enum import Enum
  import numpy
  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
label_desc: Use image node_id 3 for calling a global function with example usage: _update_confusion_matrix_variables_optimized(variables_to_update, y_true, y_pred, thresholds, multi_label, sample_weights, label_weights, thresholds_with_epsilon) and returns: ops