code: def get_segemented_ne_states(
      labels: torch.Tensor,
      predictions: torch.Tensor,
      weights: torch.Tensor,
      grouping_keys: torch.Tensor,
      eta: float,
      num_groups: int,
  ) -> Dict[str, torch.Tensor]:
      groups = torch.unique(grouping_keys)
      cross_entropy, weighted_num_samples, pos_labels, neg_labels = (
          torch.zeros(num_groups).to(labels.device),
          torch.zeros(num_groups).to(labels.device),
          torch.zeros(num_groups).to(labels.device),
          torch.zeros(num_groups).to(labels.device),
      )
      for group in groups:
          group_mask = grouping_keys == group
  
          group_labels = labels[group_mask]
          group_predictions = predictions[group_mask]
          group_weights = weights[group_mask]
  
          ce_sum_group = torch.sum(
              compute_cross_entropy(
                  labels=group_labels,
                  predictions=group_predictions,
                  weights=group_weights,
                  eta=eta,
              ),
              dim=-1,
          )
  
          weighted_num_samples_group = torch.sum(group_weights, dim=-1)
          pos_labels_group = torch.sum(
              group_weights * group_labels, dim=-1
          )
          neg_labels_group = torch.sum(
              group_weights * (1.0 - group_labels), dim=-1
          )
  
          cross_entropy[group] = ce_sum_group.item()
          weighted_num_samples[
              group
          ] = weighted_num_samples_group.item()
          pos_labels[group] = pos_labels_group.item()
          neg_labels[group] = neg_labels_group.item()
  
      # tensor size for each value is (num_groups)
      return {
          "cross_entropy_sum": cross_entropy,
          "weighted_num_samples": weighted_num_samples,
          "pos_labels": pos_labels,
          "neg_labels": neg_labels,
      }
  
imports: from typing import Any, 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_segemented_ne_states(labels, predictions, weights, grouping_keys, eta, num_groups) and returns: dict