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: Perform exponential moving average for a single step.
  
  Args:
      x: MultiHeadDampedEMA input sequences. (B, D, 1)
      state: MultiHeadDampedEMA state. (B, D, num_heads)
  
  Returns:
      out: MultiHeadDamped output sequences. (B, 1, D)
      new_state: MultiHeadDampedEMA state. (B, D, num_heads)
code: def ema_one_step(
      self, x: torch.Tensor, state: Optional[torch.Tensor] = None
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """Perform exponential moving average for a single step.
  
      Args:
          x: MultiHeadDampedEMA input sequences. (B, D, 1)
          state: MultiHeadDampedEMA state. (B, D, num_heads)
  
      Returns:
          out: MultiHeadDamped output sequences. (B, 1, D)
          new_state: MultiHeadDampedEMA state. (B, D, num_heads)
  
      """
      damping_factor, prev_timestep_weight = self.get_ema_coefficients()
  
      new_state = (damping_factor * self.ema_expansion_matrix).squeeze(
          -1
      ) * x
  
      if state is not None:
          new_state = (
              new_state + prev_timestep_weight.squeeze(-1) * state
          )
  
      out = torch.einsum(
          "bdn, dn -> bd",
          new_state,
          self.kernel_projection_matrix * self.scaling,
      )
  
      return out.unsqueeze(0), new_state
  
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 6 for calling the MultiHeadDampedEMA obj's underlying member method code with example usage: obj.ema_one_step(x, state) and returns: out,  new_state