docstr: Compute cross entropy and entropy for log probs and targets.
  
  Args:
      logits: [batch, length, num_classes] float array.
      targets: categorical targets [batch, length] int array.
  
  Returns:
      Tuple of scalar loss.
docstr_tok: ["Compute","cross","entropy","and","entropy","for","log","probs","and","targets",".","Args",":","logits",":","[","batch",",","length",",","num_classes","]","float","array",".","targets",":","categorical","targets","[","batch",",","length","]","int","array",".","Returns",":","Tuple","of","scalar","loss","."]
code: def _calculate_loss(
      self, logits: torch.Tensor, targets: torch.Tensor
  ) -> torch.Tensor:
      """Compute cross entropy and entropy for log probs and targets.
  
      Args:
          logits: [batch, length, num_classes] float array.
          targets: categorical targets [batch, length] int array.
  
      Returns:
          Tuple of scalar loss.
      """
      if len(logits.shape) != len(targets.shape) + 1:
          raise ValueError(
              "Incorrect shapes. Got shape %s logits and %s targets"
              % (str(logits.shape), str(targets.shape))
          )
      vocab_size = logits.shape[-1]
      confidence = 1.0 - self.config.label_smoothing
      low_confidence = (1.0 - confidence) / (vocab_size - 1)
      normalizing_constant = -(
          confidence * math.log(confidence)
          + (vocab_size - 1)
          * low_confidence
          * math.log(low_confidence + 1e-20)
      )
  
      # one hot
      soft_targets = targets[..., None] == torch.arange(
          vocab_size, device=targets.device
      ).reshape((1,) * len(targets.shape) + (-1,))
      soft_targets = torch.where(
          soft_targets,
          torch.full_like(soft_targets, confidence),
          torch.full_like(soft_targets, low_confidence),
      )
      soft_targets = soft_targets.to(torch.float32)
  
      # cross entropy
      total_loss = ZLossCrossEntropy.apply(
          logits, soft_targets, self.config.z_loss_factor
      )
      total_loss = total_loss - normalizing_constant
      total_loss = torch.mean(torch.sum(total_loss, dim=-1), dim=0)
      return total_loss
  
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch import nn
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from transformers.modeling_utils import PreTrainedModel
  from transformers.models.llama.modeling_llama import LlamaConfig, LlamaRMSNorm
  from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from colossalai.kernel.cuda_native.mha.flash_attn_2 import HAS_FLASH_ATTN
  from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
  from colossalai.moe.layers import SparseMLP
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.moe.utils import get_activation, set_moe_args
label_desc: Use image node_id 12 for calling the OpenMoeForCausalLM obj's underlying member method code with example usage: obj._calculate_loss(logits, targets) and returns: total_loss