class_docstr: WKVLinearAttention function definition.
docstr: WKVLinearAttention function forward pass.
  
  Args:
      time_decay: Channel-wise time decay vector. (D_att)
      time_first: Channel-wise time first vector. (D_att)
      key: Key tensor. (B, U, D_att)
      value: Value tensor. (B, U, D_att)
  
  Returns:
      out: Weighted Key-Value tensor. (B, U, D_att)
docstr_tok: ["WKVLinearAttention","function","forward","pass",".","Args",":","time_decay",":","Channel-wise","time","decay","vector",".","(","D_att",")","time_first",":","Channel-wise","time","first","vector",".","(","D_att",")","key",":","Key","tensor",".","(","B",",","U",",","D_att",")","value",":","Value","tensor",".","(","B",",","U",",","D_att",")","Returns",":","out",":","Weighted","Key-Value","tensor",".","(","B",",","U",",","D_att",")"]
code: def forward(
      ctx,
      time_decay: torch.Tensor,
      time_first: torch.Tensor,
      key: torch.Tensor,
      value: torch.tensor,
  ) -> torch.Tensor:
      """WKVLinearAttention function forward pass.
  
      Args:
          time_decay: Channel-wise time decay vector. (D_att)
          time_first: Channel-wise time first vector. (D_att)
          key: Key tensor. (B, U, D_att)
          value: Value tensor. (B, U, D_att)
  
      Returns:
          out: Weighted Key-Value tensor. (B, U, D_att)
  
      """
      batch, length, dim = key.size()
  
      assert length <= wkv_kernel.context_size, (
          f"Cannot process key of length {length} while context_size "
          f"is ({wkv_kernel.context_size}). Limit should be increased."
      )
  
      assert batch * dim % min(dim, 32) == 0, (
          f"batch size ({batch}) by dimension ({dim}) should be a multiple of "
          f"{min(dim, 32)}"
      )
  
      ctx.input_dtype = key.dtype
  
      time_decay = -torch.exp(time_decay.contiguous())
      time_first = time_first.contiguous()
  
      key = key.contiguous()
      value = value.contiguous()
  
      out = torch.empty_like(key, memory_format=torch.contiguous_format)
  
      wkv_kernel.forward(time_decay, time_first, key, value, out)
      ctx.save_for_backward(time_decay, time_first, key, value, out)
  
      return out
  
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 for calling the WKVLinearAttention obj's underlying member method code with example usage: obj.forward(ctx, time_decay, time_first, key, value) and returns: out