code: def rwkv_attention_forward(
      self,
      hidden: torch.Tensor,
      state: List[torch.Tensor] = None,
      use_cache=False,
  ):
      receptance, key, value, state = self.extract_key_value(
          hidden, state=state
      )
      layer_state = (
          tuple(s[:, :, self.layer_id] for s in state[2:])
          if state is not None
          else None
      )
  
      if hidden.device.type == "xpu":
          rwkv, layer_state = rwkv_linear_attention_xpu(
              self.time_decay,
              self.time_first,
              key,
              value,
              state=layer_state,
              return_state=use_cache,
          )
      else:
          from transformers.models.rwkv.modeling_rwkv import (
              rwkv_linear_attention_cpu,
          )
  
          rwkv, layer_state = rwkv_linear_attention_cpu(
              self.time_decay,
              self.time_first,
              key,
              value,
              state=layer_state,
              return_state=use_cache,
          )
  
      if layer_state is not None:
          state[2][:, :, self.layer_id] = layer_state[0]
          state[3][:, :, self.layer_id] = layer_state[1]
          state[4][:, :, self.layer_id] = layer_state[2]
  
      return self.output(receptance * rwkv), state
  
imports: import torch
  from typing import List
label_desc: Use image node_id 2 for calling a global function with example usage: rwkv_attention_forward(hidden, state, use_cache) and returns: self,  state