code: def squared_pearson_correlation(
      predictions: types.TensorType,
      labels: types.TensorType,
      weights: Optional[types.TensorType] = None,
  ) -> Tuple[types.TensorType, types.TensorType]:
      """Metric to compute the squared pearson correlation (R squared)."""
      with tf.compat.v1.variable_scope(
          "squared_pearson_correlation", [predictions, labels, weights]
      ):
          weighted_predictions = tf.multiply(predictions, weights)
          weighted_labels = tf.multiply(labels, weights)
  
          regression_error = tf.compat.v1.Variable(
              initial_value=0.0,
              dtype=tf.float64,
              trainable=False,
              collections=[
                  tf.compat.v1.GraphKeys.METRIC_VARIABLES,
                  tf.compat.v1.GraphKeys.LOCAL_VARIABLES,
              ],
              validate_shape=True,
              name="regression_error",
          )
  
          update_regression_error_op = tf.compat.v1.assign_add(
              regression_error,
              tf.reduce_sum(
                  input_tensor=tf.square(
                      tf.subtract(weighted_labels, weighted_predictions)
                  )
              ),
          )
  
          total_error = tf.compat.v1.Variable(
              initial_value=0.0,
              dtype=tf.float64,
              trainable=False,
              collections=[
                  tf.compat.v1.GraphKeys.METRIC_VARIABLES,
                  tf.compat.v1.GraphKeys.LOCAL_VARIABLES,
              ],
              validate_shape=True,
              name="total_error",
          )
  
          update_total_error_op = tf.compat.v1.assign_add(
              total_error,
              tf.reduce_sum(
                  input_tensor=tf.square(
                      tf.subtract(
                          weighted_labels,
                          tf.reduce_mean(input_tensor=weighted_labels),
                      )
                  )
              ),
          )
  
          div_op = tf.math.divide
          if hasattr(tf.math, "divide_no_nan"):
              div_op = tf.math.divide_no_nan
  
          def r_squared(_, regression_error, total_error):
              return tf.subtract(
                  tf.cast(1.0, tf.float64),
                  div_op(regression_error, total_error),
              )
  
          value_op = tf.distribute.get_replica_context().merge_call(
              r_squared, args=(regression_error, total_error)
          )
          update_op = tf.subtract(
              tf.cast(1.0, tf.float64),
              div_op(update_regression_error_op, update_total_error_op),
          )
          return value_op, update_op
  
file_path: model-analysis/tensorflow_model_analysis/post_export_metrics/metrics.py
imports: from typing import List, Optional, Tuple
  import numpy
  import tensorflow
  from tensorflow_model_analysis.api import types
label_desc: Use image node_id 2 for calling a global function with example usage: squared_pearson_correlation(predictions, labels, weights) and returns: value_op,  update_op, tf