class_docstr: This set of tests verify the computation logic of AUC in several
  corner cases that we know the computation results. The goal is to
  provide some confidence of the correctness of the math formula.
class_docstr_tok: ["This","set","of","tests","verify","the","computation","logic","of","AUC","in","several","corner","cases","that","we","know","the","computation","results",".","The","goal","is","to","provide","some","confidence","of","the","correctness","of","the","math","formula","."]
code: def _test_segemented_ne_helper(
      self,
      labels: torch.Tensor,
      predictions: torch.Tensor,
      weights: torch.Tensor,
      expected_ne: torch.Tensor,
      grouping_keys: torch.Tensor,
  ) -> None:
      num_task = labels.shape[0]
      batch_size = labels.shape[0]
      task_list = []
      inputs: Dict[
          str, Union[Dict[str, torch.Tensor], torch.Tensor]
      ] = {
          "predictions": {},
          "labels": {},
          "weights": {},
      }
      if grouping_keys is not None:
          inputs["required_inputs"] = {"grouping_keys": grouping_keys}
      for i in range(num_task):
          task_info = RecTaskInfo(
              name=f"Task:{i}",
              label_name="label",
              prediction_name="prediction",
              weight_name="weight",
          )
          task_list.append(task_info)
          # pyre-ignore
          inputs["predictions"][task_info.name] = predictions[i]
          # pyre-ignore
          inputs["labels"][task_info.name] = labels[i]
          # pyre-ignore
          inputs["weights"][task_info.name] = weights[i]
  
      ne = SegmentedNEMetric(
          world_size=1,
          my_rank=0,
          batch_size=batch_size,
          tasks=task_list,
          # pyre-ignore
          num_groups=max(2, torch.unique(grouping_keys)[-1].item() + 1),
      )
      # pyre-ignore
      ne.update(**inputs)
      actual_ne = ne.compute()
  
      for task_id, task in enumerate(task_list):
          for label in [0, 1]:
              cur_actual_ne = actual_ne[
                  f"segmented_ne-{task.name}|lifetime_segmented_ne_{label}"
              ]
              cur_expected_ne = expected_ne[task_id][label]
  
              torch.testing.assert_close(
                  cur_actual_ne,
                  cur_expected_ne,
                  atol=1e-4,
                  rtol=1e-4,
                  check_dtype=False,
                  equal_nan=True,
                  msg=f"Actual: {cur_actual_ne}, Expected: {cur_expected_ne}",
              )
  
file_path: torchrec/torchrec/metrics/tests/test_segmented_ne.py
imports: import unittest
  from typing import Dict, Iterable, Union
  import torch
  from torch import no_grad
  from torchrec.metrics.rec_metric import RecTaskInfo
  from torchrec.metrics.segmented_ne import SegmentedNEMetric
label_desc: Use image node_id 1 for calling the SegementedNEValueTest obj's underlying member method code with example usage: obj._test_segemented_ne_helper(labels, predictions, weights, expected_ne, grouping_keys) without return types