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 WKV with state (i.e.: for inference).
  
  Args:
      time_decay: Channel-wise time decay vector. (D_att)
      time_first: Channel-wise time first vector. (D_att)
      key: Key tensor. (B, 1, D_att)
      value: Value tensor. (B, 1, D_att)
      state: Decoder hidden states. [3 x (B, D_att)]
  
  Returns:
      output: Weighted Key-Value. (B, 1, D_att)
      state: Decoder hidden states. [3 x (B, 1, D_att)]
docstr_tok: ["Compute","WKV","with","state","(","i.e",".",":","for","inference",")",".","Args",":","time_decay",":","Channel-wise","time","decay","vector",".","(","D_att",")","time_first",":","Channel-wise","time","first","vector",".","(","D_att",")","key",":","Key","tensor",".","(","B",",","1",",","D_att",")","value",":","Value","tensor",".","(","B",",","1",",","D_att",")","state",":","Decoder","hidden","states",".","[","3","x","(","B",",","D_att",")","]","Returns",":","output",":","Weighted","Key-Value",".","(","B",",","1",",","D_att",")","state",":","Decoder","hidden","states",".","[","3","x","(","B",",","1",",","D_att",")","]"]
code: def wkv_linear_attention(
      self,
      time_decay: torch.Tensor,
      time_first: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      state: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
  ) -> Tuple[
      torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
  ]:
      """Compute WKV with state (i.e.: for inference).
  
      Args:
          time_decay: Channel-wise time decay vector. (D_att)
          time_first: Channel-wise time first vector. (D_att)
          key: Key tensor. (B, 1, D_att)
          value: Value tensor. (B, 1, D_att)
          state: Decoder hidden states. [3 x (B, D_att)]
  
      Returns:
          output: Weighted Key-Value. (B, 1, D_att)
          state: Decoder hidden states. [3 x (B, 1, D_att)]
  
      """
      num_state, den_state, max_state = state
  
      time_decay = -torch.exp(time_decay)
  
      max_for_output = torch.maximum(max_state, (time_first + key))
  
      e1 = torch.exp(max_state - max_for_output)
      e2 = torch.exp((time_first + key) - max_for_output)
  
      numerator = e1 * num_state + e2 * value
      denominator = e1 * den_state + e2
  
      max_for_state = torch.maximum(key, (max_state + time_decay))
  
      e1 = torch.exp((max_state + time_decay) - max_for_state)
      e2 = torch.exp(key - max_for_state)
  
      wkv = numerator / denominator
  
      state = [
          e1 * num_state + e2 * value,
          e1 * den_state + e2,
          max_for_state,
      ]
  
      return wkv, 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 3 for calling the SelfAttention obj's underlying member method code with example usage: obj.wkv_linear_attention(time_decay, time_first, key, value, state) and returns: wkv,  state