prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  [tensor] fix some unittests
  
  [tensor] fix some unittests
  
  [tensor] fix some unittests
  
  [BUG]: wrong in KL_approx , the order of two distribution is handled incorrectly
  compute_approx_kl for `NaiveExperienceMaker` maybe incorrect. 
  
  As motion in [Approximating KL Divergence](http://joschu.net/blog/kl-approx.html)
  
   
  $$ KL[q,p] = \mathbb{E}_{x\sim q}[\log\frac{q(x)}{p(x)}] $$
  
  let 
  
  $$ r = \frac{p(x)}{q(x)} $$
  
  note that, x is sample from distribution q. 
  
  Then 
  
  $$ KL_{approx}[q,p] = \mathbb{E}_{x\sim q}[-\log(r) + (r-1) ] $$
  
  ---
   
  In paper [Training language models to follow instructions with human feedback](https://arxiv.org/abs/2203.02155), object for actor , (e.i. reward of experience , ignore loss_ptx) 
  
  <img width="756" alt="image" src="https://user-images.githubusercontent.com/22851737/236610916-2f068c34-1508-438f-bd16-6fe6ed491e8c.png">
  
  
  <img width="795" alt="image" src="https://user-images.githubusercontent.com/22851737/236611424-17600f6e-7aca-4bdf-95bd-2ce4035bcd3a.png">
  
  So for computing KL, samples are sampled from actor model e.i  $\pi^{RL}_\phi$, instead of $\pi^{SFT}$
  
   KL in the object should be $KL[\pi^{RL}, \pi^{SFT}] =KL[q,p]$ , and $r$ of KL_approx should be $\frac{\pi^{SFT}(x)}{\pi^{RL}_\phi(x)}$
  
  --- 
  
  While on the `coati.models.utils.compute_approx_kl`
  
  ``` python 
      log_ratio = log_probs - log_probs_base
  ```
  
  and log_probs and log_probs_base correspond to actor_model and sft_model respectively.
  This should be modify to 
  
  ```python 
      log_ratio = log_probs_base - log_probs 
  ```
  
  
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `applications/Chat/coati/models/utils.py`
  Content:
  ```
  1 from typing import Optional, Union
  2 
  3 import loralib as lora
  4 import torch
  5 import torch.nn as nn
  6 import torch.nn.functional as F
  7 
  8 
  9 def compute_approx_kl(log_probs: torch.Tensor,
  10                       log_probs_base: torch.Tensor,
  11                       action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
  12     """
  13     Compute the approximate KL divergence between two distributions.
  14     Schulman blog: http://joschu.net/blog/kl-approx.html
  15 
  16     Args:
  17         log_probs: Log probabilities of the new distribution.
  18         log_probs_base: Log probabilities of the base distribution.
  19         action_mask: Mask for actions.
  20     """
  21 
  22     log_ratio = log_probs - log_probs_base
  23     approx_kl = (log_ratio.exp() - 1) - log_ratio
  24     if action_mask is not None:
  25         approx_kl = masked_mean(approx_kl, action_mask, dim=1)
  26         return approx_kl
  27     approx_kl = approx_kl.mean(dim=1)
  28     return approx_kl
  29 
  30 
  31 def compute_reward(r: Union[torch.Tensor, float],
  32                    kl_coef: float,
  33                    log_probs: torch.Tensor,
  34                    log_probs_base: torch.Tensor,
  35                    action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
  36     if kl_coef <= 0.0:
  37         return r
  38     kl = compute_approx_kl(log_probs, log_probs_base, action_mask=action_mask)
  39     reward = r - kl_coef * kl
  40     return reward
  41 
  42 
  43 def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
  44     log_probs = F.log_softmax(logits, dim=-1)
  45     log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))
  46     return log_probs_labels.squeeze(-1)
  47 
  48 
  49 def calc_action_log_probs(output: torch.Tensor, sequences: torch.LongTensor, num_actions: int) -> torch.Tensor:
  50     """Calculate action log probs.
  51 
  52     Args:
  53         output (torch.Tensor): Output tensor of Actor.forward.
  54         sequences (torch.LongTensor): Input sequences.
  55         num_actions (int): Number of actions.
  56 
  57     Returns:
  58         torch.Tensor: Action log probs.
  59     """
  60     logits = output['logits']
  61     log_probs = log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:])
  62     return log_probs[:, -num_actions:]
  63 
  64 
  65 def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor:
  66     tensor = tensor * mask
  67     tensor = tensor.sum(dim=dim)
  68     mask_sum = mask.sum(dim=dim)
  69     mean = tensor / (mask_sum + 1e-8)
  70     return mean
  71 
  72 
  73 def masked_normalize(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1, eps: float = 1e-8) -> torch.Tensor:
  74     tensor = tensor * mask
  75     mean = masked_mean(tensor, mask, dim=dim)
  76     mean_centered = tensor - mean
  77     var = masked_mean(mean_centered**2, mask, dim=dim)
  78     return mean_centered * var.clamp(min=eps).rsqrt()
  79 
  80 
  81 def normalize(tensor: torch.Tensor, dim: int = 0, eps: float = 1e-8) -> torch.Tensor:
  82     mean = tensor.mean(dim)
  83     mean_centered = tensor - mean
  84     var = (mean_centered**2).mean(dim)
  85     norm = mean_centered * var.clamp(min=eps).rsqrt()
  86     return norm
  87 
  88 
  89 def convert_to_lora(model: nn.Module,
  90                     input_size: int,
  91                     output_size: int,
  92                     lora_rank: int = 16,
  93                     lora_alpha: int = 1,
  94                     lora_dropout: float = 0.,
  95                     fan_in_fan_out: bool = False,
  96                     merge_weights: bool = True):
  97     if lora_rank > min(input_size, output_size):
  98         raise ValueError(f"LoRA rank {lora_rank} must be less or equal than {min(input_size, output_size)}")
  99 
  100     for name, module in model.named_modules():
  101         if isinstance(module, nn.Linear):
  102             module._modules[name] = lora.Linear(input_size,
  103                                                 output_size,
  104                                                 r=lora_rank,
  105                                                 lora_alpha=lora_alpha,
  106                                                 lora_dropout=lora_dropout,
  107                                                 fan_in_fan_out=fan_in_fan_out,
  108                                                 merge_weights=merge_weights)
  109
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/applications/Chat/coati/models/utils.py b/applications/Chat/coati/models/utils.py
  --- a/applications/Chat/coati/models/utils.py
  +++ b/applications/Chat/coati/models/utils.py
  @@ -19,7 +19,7 @@
           action_mask: Mask for actions.
       """
   
  -    log_ratio = log_probs - log_probs_base
  +    log_ratio = log_probs_base - log_probs
       approx_kl = (log_ratio.exp() - 1) - log_ratio
       if action_mask is not None:
           approx_kl = masked_mean(approx_kl, action_mask, dim=1)
  
verification_info: {"golden_diff": "diff --git a/applications/Chat/coati/models/utils.py b/applications/Chat/coati/models/utils.py\n--- a/applications/Chat/coati/models/utils.py\n+++ b/applications/Chat/coati/models/utils.py\n@@ -19,7 +19,7 @@\n         action_mask: Mask for actions.\n     \"\"\"\n \n-    log_ratio = log_probs - log_probs_base\n+    log_ratio = log_probs_base - log_probs\n     approx_kl = (log_ratio.exp() - 1) - log_ratio\n     if action_mask is not None:\n         approx_kl = masked_mean(approx_kl, action_mask, dim=1)\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[BUG]: wrong in KL_approx , the order of two distribution is handled incorrectly\ncompute_approx_kl for `NaiveExperienceMaker` maybe incorrect. \r\n\r\nAs motion in [Approximating KL Divergence](http://joschu.net/blog/kl-approx.html)\r\n\r\n \r\n$$ KL[q,p] = \\mathbb{E}_{x\\sim q}[\\log\\frac{q(x)}{p(x)}] $$\r\n\r\nlet \r\n\r\n$$ r = \\frac{p(x)}{q(x)} $$\r\n\r\nnote that, x is sample from distribution q. \r\n\r\nThen \r\n\r\n$$ KL_{approx}[q,p] = \\mathbb{E}_{x\\sim q}[-\\log(r) + (r-1) ] $$\r\n\r\n---\r\n \r\nIn paper [Training language models to follow instructions with human feedback](https://arxiv.org/abs/2203.02155), object for actor , (e.i. reward of experience , ignore loss_ptx) \r\n\r\n<img width=\"756\" alt=\"image\" src=\"https://user-images.githubusercontent.com/22851737/236610916-2f068c34-1508-438f-bd16-6fe6ed491e8c.png\">\r\n\r\n\r\n<img width=\"795\" alt=\"image\" src=\"https://user-images.githubusercontent.com/22851737/236611424-17600f6e-7aca-4bdf-95bd-2ce4035bcd3a.png\">\r\n\r\nSo for computing KL, samples are sampled from actor model e.i  $\\pi^{RL}_\\phi$, instead of $\\pi^{SFT}$\r\n\r\n KL in the object should be $KL[\\pi^{RL}, \\pi^{SFT}] =KL[q,p]$ , and $r$ of KL_approx should be $\\frac{\\pi^{SFT}(x)}{\\pi^{RL}_\\phi(x)}$\r\n\r\n--- \r\n\r\nWhile on the `coati.models.utils.compute_approx_kl`\r\n\r\n``` python \r\n    log_ratio = log_probs - log_probs_base\r\n```\r\n\r\nand log_probs and log_probs_base correspond to actor_model and sft_model respectively.\r\nThis should be modify to \r\n\r\n```python \r\n    log_ratio = log_probs_base - log_probs \r\n```\r\n\r\n\n", "before_files": [{"content": "from typing import Optional, Union\n\nimport loralib as lora\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef compute_approx_kl(log_probs: torch.Tensor,\n                      log_probs_base: torch.Tensor,\n                      action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:\n    \"\"\"\n    Compute the approximate KL divergence between two distributions.\n    Schulman blog: http://joschu.net/blog/kl-approx.html\n\n    Args:\n        log_probs: Log probabilities of the new distribution.\n        log_probs_base: Log probabilities of the base distribution.\n        action_mask: Mask for actions.\n    \"\"\"\n\n    log_ratio = log_probs - log_probs_base\n    approx_kl = (log_ratio.exp() - 1) - log_ratio\n    if action_mask is not None:\n        approx_kl = masked_mean(approx_kl, action_mask, dim=1)\n        return approx_kl\n    approx_kl = approx_kl.mean(dim=1)\n    return approx_kl\n\n\ndef compute_reward(r: Union[torch.Tensor, float],\n                   kl_coef: float,\n                   log_probs: torch.Tensor,\n                   log_probs_base: torch.Tensor,\n                   action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:\n    if kl_coef <= 0.0:\n        return r\n    kl = compute_approx_kl(log_probs, log_probs_base, action_mask=action_mask)\n    reward = r - kl_coef * kl\n    return reward\n\n\ndef log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n    log_probs = F.log_softmax(logits, dim=-1)\n    log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))\n    return log_probs_labels.squeeze(-1)\n\n\ndef calc_action_log_probs(output: torch.Tensor, sequences: torch.LongTensor, num_actions: int) -> torch.Tensor:\n    \"\"\"Calculate action log probs.\n\n    Args:\n        output (torch.Tensor): Output tensor of Actor.forward.\n        sequences (torch.LongTensor): Input sequences.\n        num_actions (int): Number of actions.\n\n    Returns:\n        torch.Tensor: Action log probs.\n    \"\"\"\n    logits = output['logits']\n    log_probs = log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:])\n    return log_probs[:, -num_actions:]\n\n\ndef masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor:\n    tensor = tensor * mask\n    tensor = tensor.sum(dim=dim)\n    mask_sum = mask.sum(dim=dim)\n    mean = tensor / (mask_sum + 1e-8)\n    return mean\n\n\ndef masked_normalize(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1, eps: float = 1e-8) -> torch.Tensor:\n    tensor = tensor * mask\n    mean = masked_mean(tensor, mask, dim=dim)\n    mean_centered = tensor - mean\n    var = masked_mean(mean_centered**2, mask, dim=dim)\n    return mean_centered * var.clamp(min=eps).rsqrt()\n\n\ndef normalize(tensor: torch.Tensor, dim: int = 0, eps: float = 1e-8) -> torch.Tensor:\n    mean = tensor.mean(dim)\n    mean_centered = tensor - mean\n    var = (mean_centered**2).mean(dim)\n    norm = mean_centered * var.clamp(min=eps).rsqrt()\n    return norm\n\n\ndef convert_to_lora(model: nn.Module,\n                    input_size: int,\n                    output_size: int,\n                    lora_rank: int = 16,\n                    lora_alpha: int = 1,\n                    lora_dropout: float = 0.,\n                    fan_in_fan_out: bool = False,\n                    merge_weights: bool = True):\n    if lora_rank > min(input_size, output_size):\n        raise ValueError(f\"LoRA rank {lora_rank} must be less or equal than {min(input_size, output_size)}\")\n\n    for name, module in model.named_modules():\n        if isinstance(module, nn.Linear):\n            module._modules[name] = lora.Linear(input_size,\n                                                output_size,\n                                                r=lora_rank,\n                                                lora_alpha=lora_alpha,\n                                                lora_dropout=lora_dropout,\n                                                fan_in_fan_out=fan_in_fan_out,\n                                                merge_weights=merge_weights)\n", "path": "applications/Chat/coati/models/utils.py"}], "after_files": [{"content": "from typing import Optional, Union\n\nimport loralib as lora\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef compute_approx_kl(log_probs: torch.Tensor,\n                      log_probs_base: torch.Tensor,\n                      action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:\n    \"\"\"\n    Compute the approximate KL divergence between two distributions.\n    Schulman blog: http://joschu.net/blog/kl-approx.html\n\n    Args:\n        log_probs: Log probabilities of the new distribution.\n        log_probs_base: Log probabilities of the base distribution.\n        action_mask: Mask for actions.\n    \"\"\"\n\n    log_ratio = log_probs_base - log_probs\n    approx_kl = (log_ratio.exp() - 1) - log_ratio\n    if action_mask is not None:\n        approx_kl = masked_mean(approx_kl, action_mask, dim=1)\n        return approx_kl\n    approx_kl = approx_kl.mean(dim=1)\n    return approx_kl\n\n\ndef compute_reward(r: Union[torch.Tensor, float],\n                   kl_coef: float,\n                   log_probs: torch.Tensor,\n                   log_probs_base: torch.Tensor,\n                   action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:\n    if kl_coef <= 0.0:\n        return r\n    kl = compute_approx_kl(log_probs, log_probs_base, action_mask=action_mask)\n    reward = r - kl_coef * kl\n    return reward\n\n\ndef log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n    log_probs = F.log_softmax(logits, dim=-1)\n    log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))\n    return log_probs_labels.squeeze(-1)\n\n\ndef calc_action_log_probs(output: torch.Tensor, sequences: torch.LongTensor, num_actions: int) -> torch.Tensor:\n    \"\"\"Calculate action log probs.\n\n    Args:\n        output (torch.Tensor): Output tensor of Actor.forward.\n        sequences (torch.LongTensor): Input sequences.\n        num_actions (int): Number of actions.\n\n    Returns:\n        torch.Tensor: Action log probs.\n    \"\"\"\n    logits = output['logits']\n    log_probs = log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:])\n    return log_probs[:, -num_actions:]\n\n\ndef masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor:\n    tensor = tensor * mask\n    tensor = tensor.sum(dim=dim)\n    mask_sum = mask.sum(dim=dim)\n    mean = tensor / (mask_sum + 1e-8)\n    return mean\n\n\ndef masked_normalize(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1, eps: float = 1e-8) -> torch.Tensor:\n    tensor = tensor * mask\n    mean = masked_mean(tensor, mask, dim=dim)\n    mean_centered = tensor - mean\n    var = masked_mean(mean_centered**2, mask, dim=dim)\n    return mean_centered * var.clamp(min=eps).rsqrt()\n\n\ndef normalize(tensor: torch.Tensor, dim: int = 0, eps: float = 1e-8) -> torch.Tensor:\n    mean = tensor.mean(dim)\n    mean_centered = tensor - mean\n    var = (mean_centered**2).mean(dim)\n    norm = mean_centered * var.clamp(min=eps).rsqrt()\n    return norm\n\n\ndef convert_to_lora(model: nn.Module,\n                    input_size: int,\n                    output_size: int,\n                    lora_rank: int = 16,\n                    lora_alpha: int = 1,\n                    lora_dropout: float = 0.,\n                    fan_in_fan_out: bool = False,\n                    merge_weights: bool = True):\n    if lora_rank > min(input_size, output_size):\n        raise ValueError(f\"LoRA rank {lora_rank} must be less or equal than {min(input_size, output_size)}\")\n\n    for name, module in model.named_modules():\n        if isinstance(module, nn.Linear):\n            module._modules[name] = lora.Linear(input_size,\n                                                output_size,\n                                                r=lora_rank,\n                                                lora_alpha=lora_alpha,\n                                                lora_dropout=lora_dropout,\n                                                fan_in_fan_out=fan_in_fan_out,\n                                                merge_weights=merge_weights)\n", "path": "applications/Chat/coati/models/utils.py"}]}