keras/src/metrics/regression_metrics.py/R2Score/_build
class R2Score:    def _build(self, y_true_shape, y_pred_shape):
        if len(y_pred_shape) != 2 or len(y_true_shape) != 2:
            raise ValueError(
                "R2Score expects 2D inputs with shape "
                "(batch_size, output_dim). Received input "
                f"shapes: y_pred.shape={y_pred_shape} and "
                f"y_true.shape={y_true_shape}."
            )
        if y_pred_shape[-1] is None or y_true_shape[-1] is None:
            raise ValueError(
                "R2Score expects 2D inputs with shape "
                "(batch_size, output_dim), with output_dim fully "
                "defined (not None). Received input "
                f"shapes: y_pred.shape={y_pred_shape} and "
                f"y_true.shape={y_true_shape}."
            )
        num_classes = y_pred_shape[-1]
        self.squared_sum = self.add_variable(
            name="squared_sum",
            shape=[num_classes],
            initializer=initializers.Zeros(),
        )
        self.sum = self.add_variable(
            name="sum",
            shape=[num_classes],
            initializer=initializers.Zeros(),
        )
        self.total_mse = self.add_variable(
            name="residual",
            shape=[num_classes],
            initializer=initializers.Zeros(),
        )
        self.count = self.add_variable(
            name="count",
            shape=[num_classes],
            initializer=initializers.Zeros(),
        )
        self._built = True