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: Compute time mixing.
  
  Args:
      x: SelfAttention input sequences. (B, U, size)
      state: Decoder hidden states. [5 x (B, 1, D_att, N)]
  
  Returns:
      x: SelfAttention output sequences. (B, U, size)
docstr_tok: ["Compute","time","mixing",".","Args",":","x",":","SelfAttention","input","sequences",".","(","B",",","U",",","size",")","state",":","Decoder","hidden","states",".","[","5","x","(","B",",","1",",","D_att",",","N",")","]","Returns",":","x",":","SelfAttention","output","sequences",".","(","B",",","U",",","size",")"]
code: def forward(
      self,
      x: torch.Tensor,
      state: Optional[List[torch.Tensor]] = None,
  ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]:
      """Compute time mixing.
  
      Args:
          x: SelfAttention input sequences. (B, U, size)
          state: Decoder hidden states. [5 x (B, 1, D_att, N)]
  
      Returns:
          x: SelfAttention output sequences. (B, U, size)
  
      """
      shifted_x = (
          self.time_shift(x)
          if state is None
          else state[1][..., self.block_id]
      )
  
      key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key)
      value = x * self.time_mix_value + shifted_x * (
          1 - self.time_mix_value
      )
      receptance = x * self.time_mix_receptance + shifted_x * (
          1 - self.time_mix_receptance
      )
  
      key = self.proj_key(key)
      value = self.proj_value(value)
      receptance = torch.sigmoid(self.proj_receptance(receptance))
  
      if state is not None:
          state[1][..., self.block_id] = x
  
          wkv, att_state = self.wkv_linear_attention(
              self.time_decay,
              self.time_first,
              key,
              value,
              tuple(s[..., self.block_id] for s in state[2:]),
          )
  
          state[2][..., self.block_id] = att_state[0]
          state[3][..., self.block_id] = att_state[1]
          state[4][..., self.block_id] = att_state[2]
      else:
          wkv = WKVLinearAttention.apply(
              self.time_decay, self.time_first, key, value
          )
  
      x = self.proj_output(receptance * wkv)
  
      return x, state
  
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 4 for calling the SelfAttention obj's underlying member method code with example usage: obj.forward(x, state) and returns: x,  state