def rmi_lower_bound(self, labels_4D, probs_4D):
        """
        calculate the lower bound of the region mutual information.
        Args:
                labels_4D   :   [N, C, H, W], dtype=float32
                probs_4D    :   [N, C, H, W], dtype=float32
        """
        assert labels_4D.shape == probs_4D.shape, print('shapes',
                                                        labels_4D.shape,
                                                        probs_4D.shape)

        p, s = self.rmi_pool_size, self.rmi_pool_stride
        if self.rmi_pool_stride > 1:
            if self.rmi_pool_way == 0:
                labels_4D = F.max_pool2d(labels_4D,
                                         kernel_size=p,
                                         stride=s,
                                         padding=self.kernel_padding)
                probs_4D = F.max_pool2d(probs_4D,
                                        kernel_size=p,
                                        stride=s,
                                        padding=self.kernel_padding)
            elif self.rmi_pool_way == 1:
                labels_4D = F.avg_pool2d(labels_4D,
                                         kernel_size=p,
                                         stride=s,
                                         padding=self.kernel_padding)
                probs_4D = F.avg_pool2d(probs_4D,
                                        kernel_size=p,
                                        stride=s,
                                        padding=self.kernel_padding)
            elif self.rmi_pool_way == 2:
                shape = labels_4D.shape
                new_h, new_w = shape[2] // s, shape[3] // s
                labels_4D = F.interpolate(labels_4D,
                                          size=[new_h, new_w],
                                          mode='nearest')
                probs_4D = F.interpolate(probs_4D,
                                         size=[new_h, new_w],
                                         mode='bilinear',
                                         align_corners=True)
            else:
                raise NotImplementedError("Pool way of RMI is not defined!")

        label_shape = labels_4D.shape
        n, c = label_shape[0], label_shape[1]

        la_vectors, pr_vectors = self.map_get_pairs(labels_4D,
                                                    probs_4D,
                                                    radius=self.rmi_radius,
                                                    is_combine=0)

        la_vectors = paddle.reshape(la_vectors, [n, c, self.half_d, -1])
        la_vectors = paddle.cast(la_vectors, dtype='float64')
        la_vectors.stop_gradient = True

        pr_vectors = paddle.reshape(pr_vectors, [n, c, self.half_d, -1])
        pr_vectors = paddle.cast(pr_vectors, dtype='float64')

        diag_matrix = paddle.unsqueeze(paddle.unsqueeze(paddle.eye(self.half_d),
                                                        axis=0),
                                       axis=0)
        la_vectors = la_vectors - paddle.mean(la_vectors, axis=3, keepdim=True)

        la_cov = paddle.matmul(la_vectors,
                               paddle.transpose(la_vectors, [0, 1, 3, 2]))
        pr_vectors = pr_vectors - paddle.mean(pr_vectors, axis=3, keepdim=True)
        pr_cov = paddle.matmul(pr_vectors,
                               paddle.transpose(pr_vectors, [0, 1, 3, 2]))

        pr_cov_inv = self.inverse(pr_cov +
                                  paddle.cast(diag_matrix, dtype='float64') *
                                  _POS_ALPHA)

        la_pr_cov = paddle.matmul(la_vectors,
                                  paddle.transpose(pr_vectors, [0, 1, 3, 2]))

        appro_var = la_cov - paddle.matmul(
            paddle.matmul(la_pr_cov, pr_cov_inv),
            paddle.transpose(la_pr_cov, [0, 1, 3, 2]))

        rmi_now = 0.5 * self.log_det_by_cholesky(
            appro_var + paddle.cast(diag_matrix, dtype='float64') * _POS_ALPHA)

        rmi_per_class = paddle.cast(paddle.mean(paddle.reshape(
            rmi_now, [-1, self.num_classes]),
                                                axis=0),
                                    dtype='float32')
        rmi_per_class = paddle.divide(rmi_per_class,
                                      paddle.to_tensor(float(self.half_d)))

        rmi_loss = paddle.sum(rmi_per_class) if _IS_SUM else paddle.mean(
            rmi_per_class)

        return rmi_loss

        calculate the lower bound of the region mutual information.
        Args:
                labels_4D   :   [N, C, H, W], dtype=float32
                probs_4D    :   [N, C, H, W], dtype=float32
        
paddleseg/models/losses/rmi_loss.py