class_docstr: MultiHeadDampedEMA module definition.
  
  Args:
      size: Module size.
      num_heads: Number of attention heads.
      activation: Activation function type.
      truncation_length: Maximum length for truncation.
docstr: Compute multi-dimensional damped EMA.
  
  Args:
      x: MultiHeadDampedEMA input sequence. (L, B, D)
      mask: Sequence mask. (B, 1, L)
      state: MultiHeadDampedEMA state. (B, D, num_heads)
  
  Returns:
      x: MultiHeadDampedEMA output sequence. (B, L, D)
      new_state: MultiHeadDampedEMA state. (B, D, num_heads)
code: def forward(
      self,
      x: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
      state: Optional[Dict[str, torch.Tensor]] = None,
  ) -> Union[torch.Tensor, Optional[torch.Tensor]]:
      """Compute multi-dimensional damped EMA.
  
      Args:
          x: MultiHeadDampedEMA input sequence. (L, B, D)
          mask: Sequence mask. (B, 1, L)
          state: MultiHeadDampedEMA state. (B, D, num_heads)
  
      Returns:
          x: MultiHeadDampedEMA output sequence. (B, L, D)
          new_state: MultiHeadDampedEMA state. (B, D, num_heads)
  
      """
      length = x.size(0)
  
      residual = x * self.residual_weight
  
      x = x.permute(1, 2, 0)
  
      if mask is not None:
          x = x.masked_fill(mask, 0.0)
  
      if state is not None:
          ema_output, new_state = self.ema_one_step(
              x, state=state["ema_state"]
          )
          ema_output = self.activation(ema_output + residual)
  
          return ema_output, new_state
  
      kernel = self.compute_ema_kernel(
          length
          if self.truncation_length is None
          else min(self.truncation_length, length)
      )
  
      input_fft = torch.fft.rfft(x.float(), n=(2 * length))
      kernel_fft = torch.fft.rfft(kernel.float(), n=(2 * length))
  
      ema_output = torch.fft.irfft(
          (input_fft * kernel_fft), n=(2 * length)
      )[..., :length]
      ema_output = ema_output.type_as(x)
  
      ema_output = self.activation(
          ema_output.permute(2, 0, 1) + residual
      )
  
      return ema_output, None
  
file_path: espnet/espnet2/asr_transducer/decoder/modules/mega/multi_head_damped_ema.py
imports: import math
  from typing import Dict, Optional, Tuple, Union
  import torch
label_desc: Use image node_id 7 for calling the MultiHeadDampedEMA obj's underlying member method code with example usage: obj.forward(x, mask, state) and returns: ema_output,  None, ema_output,  new_state