def METHOD_NAME(self, value):
    if self._validate_args:
        self._validate_sample(value)
    group_idx = self.group_idx
    if value.shape != group_idx.shape:
        raise ValueError(
            "GroupedNormalNormal.log_prob only supports values that have the same shape as group_idx."
        )
    prior_scale_sq = self.prior_scale.pow(2.0)
    obs_scale_sq_inv = self.obs_scale.pow(-2)
    obs_scale_sq_inv_sum = torch.zeros_like(self.prior_loc).scatter_add(
        0, self.group_idx, obs_scale_sq_inv
    )
    scale_ratio = prior_scale_sq * obs_scale_sq_inv_sum
    delta = value - self.prior_loc[group_idx]
    scaled_delta = delta * obs_scale_sq_inv
    scaled_delta_sum = torch.zeros_like(self.prior_loc).scatter_add(
        0, self.group_idx, scaled_delta
    )
    result1 = -(self.num_data_per_batch * LOG_ROOT_TWO_PI).sum()
    result2 = -0.5 * torch.log1p(scale_ratio).sum() - self.obs_scale.log().sum()
    result3 = -0.5 * torch.dot(delta, scaled_delta)
    numerator = prior_scale_sq * scaled_delta_sum.pow(2)
    result4 = 0.5 * (numerator / (1.0 + scale_ratio)).sum()
    return result1 + result2 + result3 + result4