Project: 021_Copyright-Protected_Language_Generation_via_Adapti | Function: _optimize_grid (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Grid search optimization algorithm that finds optimal combination weights by evaluating the objective function over a discretized parameter space

**KEY TERMS:** {
  "combined_range": "discretized grid for αt, βt parameters",
  "b1, b2": "αt, βt optimization variables",
  "min_idx": "optimal parameter selection"
}

**CONTEXT CODE:**             grid_size (int): Grid size for optimization.

        Returns:
            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Optimal b0, b1, b2 weights.
        """
        raise NotImplementedError("Function implementation has been masked for testing")

    def _get_logits(
        self,
        b0: torch.Tensor,
        b1: torch.Tensor,
        b2: torch.Tensor,
        logits1: torch.Tensor,
        logits2: torch.Tensor,
    ) -> torch.Tensor:
        """Compute combined logits.

        Args:
            b0 (torch.Tensor): Weight scalar b0.
            b1 (torch.Tensor): Weight scalar b1.
            b2 (torch.Tensor): Weight scalar b2.
            logits1 (torch.Tensor): Logits from model1.
            logits2 (torch.Tensor): Logits from model2.

        Returns:
            torch.Tensor: Combined logits.
        """
        # Check if we're in grid search (b1 has more than 2 dimensions) or main loop
        if b1.dim() > 2:
            # Grid search case
            # Ensure logits1 and logits2 have dimensions [batch_size, 1, vocab_size]
            if logits1.dim() == 2:
                logits1 = logits1.unsqueeze(1)  # [batch_size, 1, vocab_size]
                logits2 = logits2.unsqueeze(1)  # [batch_size, 1, vocab_size]
            # b0 has shape [batch_size, 1, 1], expand if necessary
            if b0.dim() == 2:
                b0 = b0.unsqueeze(1)  # [batch_size, 1, 1]
            combined_logits = b0 + b1 * logits1 + b2 * logits2  # [batch_size, grid_size^2, vocab_size]
        else:
            # Main loop case
            # Ensure b0, b1, b2 have shape [batch_size]
            b0 = b0.squeeze(-1)  # [batch_size]
            b1 = b1.squeeze(-1)  # [batch_size]
            b2 = b2.squeeze(-1)  # [batch_size]
            combined_logits = (
                b0.unsqueeze(-1) + b1.unsqueeze(-1) * logits1 + b2.unsqueeze(-1) * logits2
            )  # [batch_size, vocab_size]
        # combined_logits = combined_logits - torch.logsumexp(combined_logits, dim=-1, keepdim=True)
        return combined_logits

    def _optimize_grid(
        self,
        logits1: torch.Tensor,
        logits2: torch.Tensor,
        path_logprob1: torch.Tensor,
        path_logprob2: torch.Tensor,
        grid_size: int,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Optimize weights over a grid to minimize loss.

        Args:
            logits1 (torch.Tensor): Logits from model1.
            logits2 (torch.Tensor): Logits from model2.
            path_logprob1 (torch.Tensor): Accumulated log probabilities for model1.
            path_logprob2 (torch.Tensor): Accumulated log probabilities for model2.
            grid_size (int): Grid size for optimization.

        Returns:
            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Optimal b0, b1, b2 weights.
        """
        raise NotImplementedError("Function implementation has been masked for testing")

    def objective(
        self,
        b0: torch.Tensor,
        b1: torch.Tensor,
        b2: torch.Tensor,
        logits1: torch.Tensor,
        logits2: torch.Tensor,
        path_logprob1: torch.Tensor,
        path_logprob2: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the objective function for optimization.

        Args:
            b0 (torch.Tensor): Weight scalar b0.
            b1 (torch.Tensor): Weight scalar b1.
            b2 (torch.Tensor): Weight scalar b2.
            logits1 (torch.Tensor): Logits from model1.
            logits2 (torch.Tensor): Logits from model2.
            path_logprob1 (torch.Tensor): Accumulated log probabilities for model1.
            path_logprob2 (torch.Tensor): Accumulated log probabilities for model2.

        Returns:
            torch.Tensor: Computed loss.
        """
        raise NotImplementedError("Function implementation has been masked for testing")

Complete the function implementation replacing "raise NotImplementedError".