Project: 021_Copyright-Protected_Language_Generation_via_Adapti | Function: solve_optimization (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:** Core optimization algorithm that solves the copyright-protecting model fusion problem by finding optimal weights for combining two models

**KEY TERMS:** {
  "b0, b1, b2": "αt, βt, γt weights in paper",
  "grid_size": "discretization parameter for optimization",
  "path_logprob1, path_logprob2": "p(y<t|x) for each model"
}

**CONTEXT CODE:** 
        del past_key_values1, past_key_values2
        torch.cuda.empty_cache()
        gc.collect()

        if logits_list is not None:
            logits = [logit.to(input_ids.device) for logit in logits_list]
            logits = tuple(logits)
        else:
            logits = None

        if return_dict_in_generate:
            return GenerateDecoderOnlyOutput(
                sequences=input_ids,
                logits=logits,
            )
        return input_ids

    def model_forward(
        self,
        model: torch.nn.Module,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        past_key_values: Optional[Tuple[torch.Tensor]] = None,
        **model_kwargs: Any,
    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
        """Perform forward pass with the model.

        Args:
            model (torch.nn.Module): The language model.
            input_ids (torch.Tensor): Input token IDs.
            attention_mask (torch.Tensor): Attention mask.
            past_key_values (Optional[Tuple[torch.Tensor]]): Past key values for caching.
            **model_kwargs: Additional keyword arguments.

        Returns:
            Tuple[torch.Tensor, Tuple[torch.Tensor]]: Normalized logits and past key values.
        """
        # Perform forward pass with the model
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            return_dict=True,
            **model_kwargs,
        )

        next_token_norm_logits = F.log_softmax(outputs.logits.clone()[:, -1, :].float(), dim=-1)
        return next_token_norm_logits, outputs.past_key_values

    def solve_optimization(
        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]:
        """Solve the optimization problem to find optimal weights.

        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 _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:


Complete the function implementation replacing "raise NotImplementedError".