code: def compute_multihead_label_loss(
      logits: jnp.ndarray,
      one_hot_targets: jnp.ndarray,
      weights: Optional[jnp.ndarray],
      class_splits: Optional[jnp.ndarray],
      label_smoothing: Optional[float] = None,
  ):
      """Computes loss for the multi-head label task.
  
      Args:
       logits: Output of model in shape [batch, length, num_classes].
       one_hot_targets: ground truth labels of the same shape as logits.
       weights: None or array of shape [batch x ...] (rank of one_hot_targets -1).
       class_splits: start indices of class splits.
       label_smoothing: float scalar to use to smooth the one-hot labels.
  
      Returns:
        Loss value.
      """
  
      if logits.shape[-1] != class_splits[-1]:
          raise AssertionError(
              "Logit dimension must be equal to number of classes"
          )
  
      logit_splits = jnp.split(logits, class_splits, axis=-1)[:-1]
      one_hot_target_splits = jnp.split(
          one_hot_targets, class_splits, axis=-1
      )[:-1]
  
      sof_ce_losses = [
          model_utils.weighted_softmax_cross_entropy(
              logits, one_hot_targets, weights, label_smoothing
          )
          for logits, one_hot_targets in zip(
              logit_splits, one_hot_target_splits
          )
      ]
      sof_ce_loss = jnp.mean(jnp.array(sof_ce_losses))
  
      return sof_ce_loss
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple, List
  from absl import logging
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import classification_metrics_function
  from scenic.model_lib.base_models.encoder_decoder_model import encoder_decoder_metrics_function
  from scenic.model_lib.base_models.multilabel_classification_model import multilabel_classification_metrics_function
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 5 for calling a global function with example usage: compute_multihead_label_loss(logits, one_hot_targets, weights, class_splits, label_smoothing) and returns: sof_ce_loss