code: def _compute_approx_kl(
      log_probs: torch.Tensor,
      log_probs_base: torch.Tensor,
      action_mask: Optional[torch.Tensor] = None,
  ) -> torch.Tensor:
      """
      Compute the approximate KL divergence between two distributions.
      Schulman blog: http://joschu.net/blog/kl-approx.html
  
      Args:
          log_probs: Log probabilities of the new distribution.
          log_probs_base: Log probabilities of the base distribution.
          action_mask: Mask for actions.
      """
  
      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)
          return approx_kl
      approx_kl = approx_kl.mean(dim=1)
      return approx_kl
  
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 1 for calling a global function with example usage: _compute_approx_kl(log_probs, log_probs_base, action_mask) and returns: approx_kl, approx_kl