class_docstr: Unconstrained Speech Enhancement and Separation (USES) Network.
  
  Reference:
      [1] W. Zhang, K. Saijo, Z.-Q., Wang, S. Watanabe, and Y. Qian,
      “Toward Universal Speech Enhancement for Diverse Input Conditions,”
      in Proc. ASRU, 2023.
  
  args:
      input_size (int): dimension of the input feature.
      output_size (int): dimension of the output.
      bottleneck_size (int): dimension of the bottleneck feature.
          Must be a multiple of `att_heads`.
      num_blocks (int): number of processing blocks.
      num_spatial_blocks (int): number of processing blocks with channel modeling.
      segment_size (int): number of frames in each non-overlapping segment.
          This is used to segment long utterances into smaller segments for
          efficient processing.
      memory_size (int): group size of global memory tokens.
          The basic use of memory tokens is to store the history information from
          previous segments.
          The memory tokens are updated by the output of the last block after
          processing each segment.
      memory_types (int): numbre of memory token groups.
          Each group corresponds to a different type of processing, i.e.,
              the first group is used for denoising without dereverberation,
              the second group is used for denoising with dereverberation,
      rnn_type (str): type of the RNN cell in the improved Transformer layer.
      hidden_size (int): hidden dimension of the RNN cell.
      att_heads (int): number of attention heads in Transformer.
      dropout (float): dropout ratio. Default is 0.
      activation (str): non-linear activation function applied in each block.
      bidirectional (bool): whether the RNN layers are bidirectional.
      norm_type (str): normalization type in the improved Transformer layer.
      ch_mode (str): mode of channel modeling.
          Select from "att" and "tac".
      ch_att_dim (int): dimension of the channel attention.
      eps (float): epsilon for layer normalization.
docstr: USES forward.
  
  Args:
      input (torch.Tensor): input feature (batch, mics, input_size, freq, time)
      ref_channel (None or int): index of the reference channel.
          if None, simply average all channels.
          if int, take the specified channel instead of averaging.
      mem_idx (None or int): index of the memory token group.
          if None, use the only group of memory tokens in the model.
          if int, use the specified group from multiple existing groups.
  Returns:
      output (torch.Tensor): output feature (batch, output_size, freq, time)
docstr_tok: ["USES","forward",".","Args",":","input","(","torch.Tensor",")",":","input","feature","(","batch",",","mics",",","input_size",",","freq",",","time",")","ref_channel","(","None","or","int",")",":","index","of","the","reference","channel",".","if","None",",","simply","average","all","channels",".","if","int",",","take","the","specified","channel","instead","of","averaging",".","mem_idx","(","None","or","int",")",":","index","of","the","memory","token","group",".","if","None",",","use","the","only","group","of","memory","tokens","in","the","model",".","if","int",",","use","the","specified","group","from","multiple","existing","groups",".","Returns",":","output","(","torch.Tensor",")",":","output","feature","(","batch",",","output_size",",","freq",",","time",")"]
code: def forward(self, input, ref_channel=None, mem_idx=None):
      """USES forward.
  
      Args:
          input (torch.Tensor): input feature (batch, mics, input_size, freq, time)
          ref_channel (None or int): index of the reference channel.
              if None, simply average all channels.
              if int, take the specified channel instead of averaging.
          mem_idx (None or int): index of the memory token group.
              if None, use the only group of memory tokens in the model.
              if int, use the specified group from multiple existing groups.
      Returns:
          output (torch.Tensor): output feature (batch, output_size, freq, time)
      """
      B, C, N, F, T = input.shape
      output = self.layer_norm(input.reshape(B * C, N, -1))
      # B, C, bn, F, T
      output = self.bottleneck_conv1x1(output).reshape(B, C, -1, F, T)
      bn = output.size(2)
  
      # Divide the input into non-overlapping segments
      num_seg, res = divmod(T, self.segment_size)
      if res > 0:
          # pad the last segment if necessary
          output = nn.functional.pad(
              output, (0, self.segment_size - res)
          )
          num_seg += 1
      if self.training and num_seg < 2:
          warnings.warn(f"The input is too short for training: {T}")
      output = output.reshape(B, C, bn, F, num_seg, self.segment_size)
  
      # Segment-by-segment processing for memory-efficient processing
      ret = []
      mem = None
      for n in range(num_seg):
          out = output[..., n, :]
          if mem is None:
              # initialize memory tokens for the first segment
              if mem_idx is not None:
                  mem = self.memory_tokens[mem_idx].repeat(
                      B, C, 1, F, 1
                  )
              elif self.memory_types > 1:
                  mem = self.memory_tokens[0].repeat(B, C, 1, F, 1)
              else:
                  mem = self.memory_tokens.repeat(B, C, 1, F, 1)
              out = torch.cat([mem, out], dim=-1)
          else:
              # reuse memory tokens from the last segment
              if mem.size(1) < C:
                  mem = mem.repeat(1, C // mem.size(1), 1, 1, 1)
              out = torch.cat([mem, out], dim=-1)
          for block in self.atf_blocks:
              out = block(out, ref_channel=ref_channel)
          mem, out = (
              out[..., : self.memory_size],
              out[..., self.memory_size :],
          )
          ret.append(out)
  
      output = torch.cat(ret, dim=-1)[..., :T]
      with torch.cuda.amp.autocast(enabled=False):
          output = self.output(output.mean(1))  # B, output_size, F, T
      return output
  
imports: import warnings
  import torch
  import torch.nn
  from espnet2.enh.layers.dptnet import ImprovedTransformerLayer
  from espnet2.enh.layers.tcn import ChannelwiseLayerNorm
  from espnet2.torch_utils.get_layer_from_string import get_layer
label_desc: Use image node_id 2 for calling the USES obj's underlying member method code with example usage: obj.forward(input, ref_channel, mem_idx) and returns: output