code: def get_ndcg_states(
      labels: torch.Tensor,
      predictions: torch.Tensor,
      weights: torch.Tensor,
      session_ids: torch.Tensor,
      exponential_gain: bool,
  ) -> Dict[str, torch.Tensor]:
      n_tasks = labels.shape[0]
      sum_ndcg = torch.zeros(n_tasks, dtype=torch.double).to(
          labels.device
      )
  
      for i in range(n_tasks):
          session_lengths = session_ids_to_lengths(
              session_ids=session_ids[i], device=predictions.device
          )
  
          sum_ndcg[i] = compute_lambda_ndcg(
              prediction=predictions[i],
              label=labels[i],
              weight=weights[i],
              session_lengths=session_lengths,
              use_exp_gain=exponential_gain,
          )
  
      return {
          SUM_NDCG: sum_ndcg,
          NUM_SESSIONS: session_lengths.shape[-1],
      }
  
imports: from typing import Any, cast, Dict, List, Optional, Type
  import torch
  from torch import distributed
  from torchrec.metrics.metrics_config import RecComputeMode, RecTaskInfo
  from torchrec.metrics.metrics_namespace import MetricName, MetricNamespace, MetricPrefix
  from torchrec.metrics.rec_metric import MetricComputationReport, RecMetric, RecMetricComputation, RecMetricException
label_desc: Use image node_id 6 for calling a global function with example usage: get_ndcg_states(labels, predictions, weights, session_ids, exponential_gain) and returns: dict