code: def compute_reward(
      r: Union[torch.Tensor, float],
      kl_coef: float,
      log_probs: torch.Tensor,
      log_probs_base: torch.Tensor,
      action_mask: Optional[torch.Tensor] = None,
  ) -> torch.Tensor:
      if kl_coef <= 0.0:
          return r
      kl = _compute_approx_kl(
          log_probs, log_probs_base, action_mask=action_mask
      )
      reward = r - kl_coef * kl
      return reward
  
file_path: ColossalAI/applications/Chat/coati/models/utils.py
imports: from typing import Optional, Union
  import torch
  import torch.nn.functional
label_desc: Use image node_id 2 for calling a global function with example usage: compute_reward(r, kl_coef, log_probs, log_probs_base, action_mask) and returns: reward, r