code: def _ref_metric(
      preds: Tensor,
      target: Tensor,
      scale_invariant: bool,
      zero_mean: bool,
  ):
      # According to the original paper, the sa-sdr equals to si-sdr with inputs concatenated over the speaker
      # dimension if scale_invariant==True. Accordingly, for scale_invariant==False, the sa-sdr equals to snr.
      # shape: preds [BATCH_SIZE, Spk, Time] , target [BATCH_SIZE, Spk, Time]
      # or shape: preds [NUM_BATCHES*BATCH_SIZE, Spk, Time], target [NUM_BATCHES*BATCH_SIZE, Spk, Time]
  
      if zero_mean:
          target = target - torch.mean(target, dim=-1, keepdim=True)
          preds = preds - torch.mean(preds, dim=-1, keepdim=True)
  
      preds = preds.reshape(
          preds.shape[0], preds.shape[1] * preds.shape[2]
      )
      target = target.reshape(
          target.shape[0], target.shape[1] * target.shape[2]
      )
      if scale_invariant:
          return scale_invariant_signal_distortion_ratio(
              preds=preds, target=target, zero_mean=False
          )
      return signal_noise_ratio(
          preds=preds, target=target, zero_mean=zero_mean
      )
  
file_path: torchmetrics/tests/unittests/audio/test_sa_sdr.py
imports: from functools import partial
  import pytest
  import torch
  from torch import Tensor
  from torchmetrics.audio import SourceAggregatedSignalDistortionRatio
  from torchmetrics.functional.audio import scale_invariant_signal_distortion_ratio, signal_noise_ratio, source_aggregated_signal_distortion_ratio
  from unittests import BATCH_SIZE, NUM_BATCHES, _Input
  from unittests.helpers import seed_all
  from unittests.helpers.testers import MetricTester
label_desc: Use image node_id 1 for calling a global function with example usage: _ref_metric(preds, target, scale_invariant, zero_mean) and returns: signal_noise_ratio, scale_invariant_signal_distortion_ratio