code: def get_ne_states(
      labels: torch.Tensor,
      predictions: torch.Tensor,
      weights: torch.Tensor,
      eta: float,
  ) -> Dict[str, torch.Tensor]:
      cross_entropy = compute_cross_entropy(
          labels,
          predictions,
          weights,
          eta,
      )
      return {
          "cross_entropy_sum": torch.sum(cross_entropy, dim=-1),
          "weighted_num_samples": torch.sum(weights, dim=-1),
          "pos_labels": torch.sum(weights * labels, dim=-1),
          "neg_labels": torch.sum(weights * (1.0 - labels), dim=-1),
      }
  
file_path: torchrec/torchrec/metrics/ne.py
imports: from typing import Any, cast, Dict, List, Optional, Type
  import torch
  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 5 for calling a global function with example usage: get_ne_states(labels, predictions, weights, eta) and returns: dict