code: def generate_model_outputs_cases() -> (
      Iterable[Dict[str, torch._tensor.Tensor]]
  ):
      return [
          # base condition
          {
              "labels": torch.tensor([[1, 0, 0, 1, 1]]),
              "predictions": torch.tensor([[0.2, 0.6, 0.8, 0.4, 0.9]]),
              "weights": torch.tensor([[0.13, 0.2, 0.5, 0.8, 0.75]]),
              "grouping_keys": torch.tensor([0, 1, 0, 1, 1]),
              "expected_ne": torch.tensor([[3.1615, 1.6004]]),
          },
          # one sided, edge case 1s
          {
              "labels": torch.tensor([[1, 0, 0, 1, 1]]),
              "predictions": torch.tensor([[0.2, 0.6, 0.8, 0.4, 0.9]]),
              "weights": torch.tensor([[0.13, 0.2, 0.5, 0.8, 0.75]]),
              "grouping_keys": torch.tensor([1, 1, 1, 1, 1]),
              "expected_ne": torch.tensor([[torch.nan, 1.3936]]),
          },
          # one sided, edge case 0s
          {
              "labels": torch.tensor([[1, 0, 0, 1, 1]]),
              "predictions": torch.tensor([[0.2, 0.6, 0.8, 0.4, 0.9]]),
              "weights": torch.tensor([[0.13, 0.2, 0.5, 0.8, 0.75]]),
              "grouping_keys": torch.tensor([0, 0, 0, 0, 0]),
              "expected_ne": torch.tensor([[1.3936, torch.nan]]),
          },
          # three labels,
          {
              "labels": torch.tensor([[1, 0, 0, 1, 1, 0]]),
              "predictions": torch.tensor(
                  [[0.2, 0.6, 0.8, 0.4, 0.9, 0.4]]
              ),
              "weights": torch.tensor(
                  [[0.13, 0.2, 0.5, 0.8, 0.75, 0.4]]
              ),
              "grouping_keys": torch.tensor([0, 1, 0, 1, 2, 2]),
              "expected_ne": torch.tensor([[3.1615, 1.8311, 0.3814]]),
          },
          # two tasks
          {
              "labels": torch.tensor(
                  [[1, 0, 0, 1, 1], [1, 0, 0, 1, 1]]
              ),
              "predictions": torch.tensor(
                  [
                      [0.2, 0.6, 0.8, 0.4, 0.9],
                      [0.6, 0.2, 0.4, 0.8, 0.9],
                  ]
              ),
              "weights": torch.tensor(
                  [
                      [0.13, 0.2, 0.5, 0.8, 0.75],
                      [0.13, 0.2, 0.5, 0.8, 0.75],
                  ]
              ),
              "grouping_keys": torch.tensor(
                  [0, 1, 0, 1, 1]
              ),  # for this case, both tasks have same groupings
              "expected_ne": torch.tensor(
                  [[3.1615, 1.6004], [1.0034, 0.4859]]
              ),
          },
      ]
  
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 a global function with example usage: generate_model_outputs_cases() and returns: list