Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/metrics/regression/explained_variance.py (Python)

Code before:
        super().__init__(
            compute_on_step=compute_on_step,
            ddp_sync_on_step=ddp_sync_on_step,
            process_group=process_group,
        )
        allowed_multioutput = ('raw_values', 'uniform_average', 'variance_weighted')
        assert multioutput in allowed_multioutput, \
            'Invalid input to argument `multioutput`. Choose one of the following:' \
            f'{allowed_multioutput}'
        self.multioutput = multioutput
        self.add_state("y", default=[], dist_reduce_fx=None)
        self.add_state("y_pred", default=[], dist_reduce_fx=None)

        rank_zero_warn('Metric `ExplainedVariance` will save all targets and'
                       ' predictions in buffer. For large datasets this may lead'
                       ' to large memory footprint.')

    def update(self, preds: torch.Tensor, target: torch.Tensor):
        """
        Update state with predictions and targets.

        Args:
            preds: Predictions from model
            target: Ground truth values
        """
        assert preds.shape == target.shape, \
            'Predictions and targets are expected to have the same shape'
        self.y.append(target)
        self.y_pred.append(preds)

    def compute(self):
        """
        Computes explained variance over state.
        """
        y_true = torch.cat(self.y, dim=0)
        y_pred = torch.cat(self.y_pred, dim=0)

        y_diff_avg = torch.mean(y_true - y_pred, dim=0)
        numerator = torch.mean((y_true - y_pred - y_diff_avg) ** 2, dim=0)

        y_true_avg = torch.mean(y_true, dim=0)
        denominator = torch.mean((y_true - y_true_avg) ** 2, dim=0)

        # Take care of division by zero
        nonzero_numerator = numerator != 0
        nonzero_denominator = denominator != 0
        valid_score = nonzero_numerator & nonzero_denominator
        output_scores = torch.ones_like(y_diff_avg)
        output_scores[valid_score] = 1.0 - (numerator[valid_score] / denominator[valid_score])
        output_scores[nonzero_numerator & ~nonzero_denominator] = 0.


Review comment (bug):
lets consider using Error clases instead...
https://medium.com/@jadhavmanoj/python-what-is-raise-and-assert-statement-c3908697bc62#:~:text=Python's%20assert%20statement%20is%20a,a%20bug%20in%20your%20program.

Code after:
        super().__init__(
            compute_on_step=compute_on_step,
            ddp_sync_on_step=ddp_sync_on_step,
            process_group=process_group,
        )
        allowed_multioutput = ('raw_values', 'uniform_average', 'variance_weighted')
        if multioutput not in allowed_multioutput:
            raise ValueError(
                f'Invalid input to argument `multioutput`. Choose one of the following: {allowed_multioutput}'
            )
        self.multioutput = multioutput
        self.add_state("y", default=[], dist_reduce_fx=None)
        self.add_state("y_pred", default=[], dist_reduce_fx=None)

        rank_zero_warn('Metric `ExplainedVariance` will save all targets and'
                       ' predictions in buffer. For large datasets this may lead'
                       ' to large memory footprint.')

    def update(self, preds: torch.Tensor, target: torch.Tensor):
        """
        Update state with predictions and targets.

        Args:
            preds: Predictions from model
            target: Ground truth values
        """
        self._check_same_shape(preds, target)
        self.y.append(target)
        self.y_pred.append(preds)

    def compute(self):
        """
        Computes explained variance over state.
        """
        y_true = torch.cat(self.y, dim=0)
        y_pred = torch.cat(self.y_pred, dim=0)

        y_diff_avg = torch.mean(y_true - y_pred, dim=0)
        numerator = torch.mean((y_true - y_pred - y_diff_avg) ** 2, dim=0)

        y_true_avg = torch.mean(y_true, dim=0)
        denominator = torch.mean((y_true - y_true_avg) ** 2, dim=0)

        # Take care of division by zero
        nonzero_numerator = numerator != 0
        nonzero_denominator = denominator != 0
        valid_score = nonzero_numerator & nonzero_denominator
        output_scores = torch.ones_like(y_diff_avg)
        output_scores[valid_score] = 1.0 - (numerator[valid_score] / denominator[valid_score])
        output_scores[nonzero_numerator & ~nonzero_denominator] = 0.