docstr: Time-Domain SpeakerBeam Extractor.
  
  Args:
      input_dim: input feature dimension
      layer: int, number of layers in each stack
      stack: int, number of stacks
      bottleneck_dim: bottleneck dimension
      hidden_dim: number of convolution channel
      skip_dim: int, number of skip connection channels
      kernel: int, kernel size.
      causal: bool, defalut False.
      norm_type: str, choose from 'BN', 'gLN', 'cLN'
      pre_nonlinear: the nonlinear function right before mask estimation
                 select from 'prelu', 'relu', 'tanh', 'sigmoid', 'linear'
      nonlinear: the nonlinear function for mask estimation,
                 select from 'relu', 'tanh', 'sigmoid', 'linear'
      i_adapt_layer: int, index of adaptation layer
      adapt_layer_type: str, type of adaptation layer
          see espnet2.enh.layers.adapt_layers for options
      adapt_enroll_dim: int, dimensionality of the speaker embedding
      use_spk_emb: bool, whether to use speaker embeddings as enrollment
      spk_emb_dim: int, dimension of input speaker embeddings
                   only used when `use_spk_emb` is True
code: def __init__(
      self,
      input_dim: int,
      layer: int = 8,
      stack: int = 3,
      bottleneck_dim: int = 128,
      hidden_dim: int = 512,
      skip_dim: int = 128,
      kernel: int = 3,
      causal: bool = False,
      norm_type: str = "gLN",
      pre_nonlinear: str = "prelu",
      nonlinear: str = "relu",
      # enrollment related arguments
      i_adapt_layer: int = 7,
      adapt_layer_type: str = "mul",
      adapt_enroll_dim: int = 128,
      use_spk_emb: bool = False,
      spk_emb_dim: int = 256,
  ):
      """Time-Domain SpeakerBeam Extractor.
  
      Args:
          input_dim: input feature dimension
          layer: int, number of layers in each stack
          stack: int, number of stacks
          bottleneck_dim: bottleneck dimension
          hidden_dim: number of convolution channel
          skip_dim: int, number of skip connection channels
          kernel: int, kernel size.
          causal: bool, defalut False.
          norm_type: str, choose from 'BN', 'gLN', 'cLN'
          pre_nonlinear: the nonlinear function right before mask estimation
                     select from 'prelu', 'relu', 'tanh', 'sigmoid', 'linear'
          nonlinear: the nonlinear function for mask estimation,
                     select from 'relu', 'tanh', 'sigmoid', 'linear'
          i_adapt_layer: int, index of adaptation layer
          adapt_layer_type: str, type of adaptation layer
              see espnet2.enh.layers.adapt_layers for options
          adapt_enroll_dim: int, dimensionality of the speaker embedding
          use_spk_emb: bool, whether to use speaker embeddings as enrollment
          spk_emb_dim: int, dimension of input speaker embeddings
                       only used when `use_spk_emb` is True
      """
      super().__init__()
  
      if pre_nonlinear not in (
          "sigmoid",
          "prelu",
          "relu",
          "tanh",
          "linear",
      ):
          raise ValueError(
              "Not supporting pre_nonlinear={}".format(pre_nonlinear)
          )
      if nonlinear not in ("sigmoid", "relu", "tanh", "linear"):
          raise ValueError(
              "Not supporting nonlinear={}".format(nonlinear)
          )
  
      self.tcn = TemporalConvNetInformed(
          N=input_dim,
          B=bottleneck_dim,
          H=hidden_dim,
          P=kernel,
          X=layer,
          R=stack,
          Sc=skip_dim,
          out_channel=None,
          norm_type=norm_type,
          causal=causal,
          pre_mask_nonlinear=pre_nonlinear,
          mask_nonlinear=nonlinear,
          i_adapt_layer=i_adapt_layer,
          adapt_layer_type=adapt_layer_type,
          adapt_enroll_dim=adapt_enroll_dim,
      )
  
      # Auxiliary network
      self.use_spk_emb = use_spk_emb
      if use_spk_emb:
          self.auxiliary_net = torch.nn.Conv1d(
              spk_emb_dim,
              adapt_enroll_dim
              if skip_dim is None
              else adapt_enroll_dim * 2,
              1,
          )
      else:
          self.auxiliary_net = TemporalConvNet(
              N=input_dim,
              B=bottleneck_dim,
              H=hidden_dim,
              P=kernel,
              X=layer,
              R=1,
              C=1,
              Sc=skip_dim,
              out_channel=adapt_enroll_dim
              if skip_dim is None
              else adapt_enroll_dim * 2,
              norm_type=norm_type,
              causal=False,
              pre_mask_nonlinear=pre_nonlinear,
              mask_nonlinear="linear",
          )
  
file_path: espnet/espnet2/enh/extractor/td_speakerbeam_extractor.py
imports: from collections import OrderedDict
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  from torch_complex.tensor import ComplexTensor
  from espnet2.enh.extractor.abs_extractor import AbsExtractor
  from espnet2.enh.layers.complex_utils import is_complex
  from espnet2.enh.layers.tcn import TemporalConvNet, TemporalConvNetInformed
  from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
label_desc: Use image node_id 1 to create a new TDSpeakerBeamExtractor object from inherited base classes: AbsExtractor with example: obj = TDSpeakerBeamExtractor(input_dim, layer, stack, bottleneck_dim, hidden_dim, skip_dim, kernel, causal, norm_type, pre_nonlinear, nonlinear, i_adapt_layer, adapt_layer_type, adapt_enroll_dim, use_spk_emb, spk_emb_dim)