class_docstr: Parametric multichannel Wiener filter, with an adjustable
  tradeoff between noise reduction and speech distortion.
  It supports automatic reference channel selection based
  on the estimated output SNR.
  
  Args:
      beta: Parameter of the parameteric filter, tradeoff between noise reduction
            and speech distortion (0: MVDR, 1: MWF).
      rank: Rank assumption for the speech covariance matrix.
      postfilter: Optional postfilter. If None, no postfilter is applied.
      ref_channel: Optional, reference channel. If None, it will be estimated automatically.
      ref_hard: If true, estimate a hard (one-hot) reference. If false, a soft reference.
      ref_hard_use_grad: If true, use straight-through gradient when using the hard reference
      ref_subband_weighting: If true, use subband weighting when estimating reference channel
      num_subbands: Optional, used to determine the parameter size for reference estimation
      diag_reg: Optional, diagonal regularization for the multichannel filter
      eps: Small regularization constant to avoid division by zero
  
  References:
      - Souden et al, On Optimal Frequency-Domain Multichannel Linear Filtering for Noise Reduction, 2010
docstr: Apply diagonal regularization on psd.
  
  Args:
      psd: tensor, shape (..., C, C)
  
  Returns:
      Tensor, same shape as input.
docstr_tok: ["Apply","diagonal","regularization","on","psd",".","Args",":","psd",":","tensor",",","shape","(","...",",","C",",","C",")","Returns",":","Tensor",",","same","shape","as","input","."]
code: def apply_diag_reg(self, psd: torch.Tensor) -> torch.Tensor:
      """Apply diagonal regularization on psd.
  
      Args:
          psd: tensor, shape (..., C, C)
  
      Returns:
          Tensor, same shape as input.
      """
      # Regularization: diag_reg * trace(psd) + eps
      diag_reg = self.diag_reg * self.trace(psd).real + self.eps
  
      # Apply regularization
      psd = psd + torch.diag_embed(
          diag_reg.unsqueeze(-1)
          * torch.ones(psd.shape[-1], device=psd.device)
      )
  
      return psd
  
imports: import random
  from typing import Callable, Optional
  import torch
  from nemo.collections.asr.parts.submodules.multi_head_attention import MultiHeadAttention
  from nemo.core.classes import NeuralModule, typecheck
  from nemo.core.neural_types import AudioSignal, FloatType, NeuralType, SpectrogramType
  from nemo.utils import logging
label_desc: Use image node_id 3 for calling the ParametricMultichannelWienerFilter obj's underlying member method code with example usage: obj.apply_diag_reg(psd) and returns: psd