code: def groupby_apply(
      keys: torch.Tensor,
      values: torch.Tensor,
      bins: int = 95,
      reduction: str = "mean",
      return_histogram: bool = False,
  ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
      """
      Groupby apply for torch tensors
  
      Args:
          keys: tensor of groups (``0`` to ``bins``)
          values: values to aggregate - same size as keys
          bins: total number of groups
          reduction: either "mean" or "sum"
          return_histogram: if to return histogram on top
  
      Returns:
          tensor of size ``bins`` with aggregated values and optionally with counts of values
      """
      if reduction == "mean":
          reduce = torch.mean
      elif reduction == "sum":
          reduce = torch.sum
      else:
          raise ValueError(f"Unknown reduction '{reduction}'")
      uniques, counts = keys.unique(return_counts=True)
      groups = torch.stack(
          [
              reduce(item)
              for item in torch.split_with_sizes(values, tuple(counts))
          ]
      )
      reduced = torch.zeros(
          bins, dtype=values.dtype, device=values.device
      ).scatter(dim=0, index=uniques, src=groups)
      if return_histogram:
          hist = torch.zeros(
              bins, dtype=torch.long, device=values.device
          ).scatter(dim=0, index=uniques, src=counts)
          return reduced, hist
      else:
          return reduced
  
imports: from collections import namedtuple
  from contextlib import redirect_stdout
  import inspect
  import os
  from typing import Any, Callable, Dict, List, Tuple, Union
  import lightning.pytorch
  import torch
  from torch import nn
  from torch.fft import irfft, rfft
  import torch.nn.functional
  from torch.nn.utils import rnn
label_desc: Use image node_id 2 for calling a global function with example usage: groupby_apply(keys, values, bins, reduction, return_histogram) and returns: reduced,  hist, reduced