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: Construct a SelfAttention object.
docstr_tok: ["Construct","a","SelfAttention","object","."]
code: def __init__(
      self,
      size: int,
      attention_size: int,
      context_size: int,
      block_id: int,
      num_blocks: int,
  ) -> None:
      """Construct a SelfAttention object."""
      super().__init__()
  
      load_wkv_kernel(context_size)
  
      self.time_shift = torch.nn.ZeroPad2d((0, 0, 1, -1))
  
      self.time_decay = torch.nn.Parameter(torch.empty(attention_size))
      self.time_first = torch.nn.Parameter(torch.empty(attention_size))
  
      self.time_mix_key = torch.nn.Parameter(torch.empty(1, 1, size))
      self.time_mix_value = torch.nn.Parameter(torch.empty(1, 1, size))
      self.time_mix_receptance = torch.nn.Parameter(
          torch.empty(1, 1, size)
      )
  
      self.proj_key = torch.nn.Linear(size, attention_size, bias=True)
      self.proj_value = torch.nn.Linear(size, attention_size, bias=True)
      self.proj_receptance = torch.nn.Linear(
          size, attention_size, bias=True
      )
  
      self.proj_output = torch.nn.Linear(
          attention_size, size, bias=True
      )
  
      self.block_id = block_id
  
      self.reset_parameters(size, attention_size, block_id, num_blocks)
  
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 1 to create a new SelfAttention object from inherited base classes: torch.nn with example: obj = SelfAttention(size, attention_size, context_size, block_id, num_blocks)