code: def compute_auprc_per_group(
      n_tasks: int,
      predictions: torch.Tensor,
      labels: torch.Tensor,
      weights: torch.Tensor,
      grouping_keys: torch.Tensor,
  ) -> torch.Tensor:
      """
      Computes AUPRC (Area Under the Curve) for binary classification for groups of predictions/labels.
      Args:
          n_tasks (int): number of tasks
          predictions (torch.Tensor): tensor of size (n_tasks, n_examples)
          labels (torch.Tensor): tensor of size (n_tasks, n_examples)
          weights (torch.Tensor): tensor of size (n_tasks, n_examples)
          grouping_keys (torch.Tensor): tensor of size (n_examples,)
  
      Returns:
          torch.Tensor: tensor of size (n_tasks,), average of AUPRCs per group.
      """
      auprcs = []
      if grouping_keys.numel() != 0 and grouping_keys[0] == -1:
          # we added padding  as the first elements during init to avoid floating point exception in sync()
          # removing the paddings to avoid numerical errors.
          grouping_keys = grouping_keys[1:]
          predictions = predictions[:, 1:]
          labels = labels[:, 1:]
          weights = weights[:, 1:]
  
      # get unique group indices
      group_indices = torch.unique(grouping_keys)
  
      for predictions_i, labels_i, weights_i in zip(
          predictions, labels, weights
      ):
          # Loop over each group
          auprc_groups_sum = torch.tensor([0], dtype=torch.float32)
          for group_idx in group_indices:
              # get predictions, labels, and weights for this group
              group_mask = grouping_keys == group_idx
              grouped_predictions = predictions_i[group_mask]
              grouped_labels = labels_i[group_mask]
              grouped_weights = weights_i[group_mask]
  
              auprc = _compute_auprc_helper(
                  grouped_predictions, grouped_labels, grouped_weights
              )
              auprc_groups_sum = auprc_groups_sum.to(auprc.device)
              auprc_groups_sum += auprc.view(1)
          avg_auprc = (
              auprc_groups_sum / len(group_indices)
              if len(group_indices) > 0
              else torch.tensor([0.5], dtype=torch.float32)
          )
          auprcs.append(avg_auprc)
      return torch.cat(auprcs)
  
imports: from functools import partial
  from typing import Any, cast, Dict, List, Optional, Type
  import torch
  import torch.distributed
  import torch.nn.functional
  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 4 for calling a global function with example usage: compute_auprc_per_group(n_tasks, predictions, labels, weights, grouping_keys) and returns: torch