class_docstr: SelfAttention module definition.
  
  Args:
      size: Input/Output size.
      attention_size: Attention hidden size.
      context_size: Context size for WKV kernel.
      block_id: Block index.
      num_blocks: Number of blocks in the architecture.
docstr: Reset module parameters.
  
  Args:
      size: Block size.
      attention_size: Attention hidden size.
      block_id: Block index.
      num_blocks: Number of blocks in the architecture.
docstr_tok: ["Reset","module","parameters",".","Args",":","size",":","Block","size",".","attention_size",":","Attention","hidden","size",".","block_id",":","Block","index",".","num_blocks",":","Number","of","blocks","in","the","architecture","."]
code: def reset_parameters(
      self,
      size: int,
      attention_size: int,
      block_id: int,
      num_blocks: int,
  ) -> None:
      """Reset module parameters.
  
      Args:
          size: Block size.
          attention_size: Attention hidden size.
          block_id: Block index.
          num_blocks: Number of blocks in the architecture.
  
      """
      ratio_0_to_1 = block_id / (num_blocks - 1)
      ratio_1_to_almost0 = 1.0 - (block_id / num_blocks)
  
      time_weight = torch.ones(1, 1, size)
  
      for i in range(size):
          time_weight[0, 0, i] = i / size
  
      decay_speed = [
          -5
          + 8 * (h / (attention_size - 1)) ** (0.7 + 1.3 * ratio_0_to_1)
          for h in range(attention_size)
      ]
      decay_speed = torch.tensor(
          decay_speed,
          dtype=self.time_decay.dtype,
          device=self.time_decay.device,
      )
  
      zigzag = (
          torch.tensor(
              [(i + 1) % 3 - 1 for i in range(attention_size)],
              dtype=self.time_first.dtype,
              device=self.time_first.device,
          )
          * 0.5
      )
  
      with torch.no_grad():
          self.time_decay.data = decay_speed
          self.time_first.data = torch.ones_like(
              self.time_first * math.log(0.3) + zigzag
          )
  
          self.time_mix_key.data = torch.pow(
              time_weight, ratio_1_to_almost0
          )
          self.time_mix_value.data = (
              torch.pow(time_weight, ratio_1_to_almost0)
              + 0.3 * ratio_0_to_1
          )
          self.time_mix_receptance.data = torch.pow(
              time_weight, 0.5 * ratio_1_to_almost0
          )
  
file_path: espnet/espnet2/asr_transducer/decoder/modules/rwkv/attention.py
imports: import math
  from importlib.util import find_spec
  from pathlib import Path
  from typing import List, Optional, Tuple, Union
  import torch
label_desc: Use image node_id 2 for calling the SelfAttention obj's underlying member method code with example usage: obj.reset_parameters(size, attention_size, block_id, num_blocks) without return types