code: def rwkv_linear_attention_cpu(
      time_decay, time_first, key, value, state=None, return_state=False
  ):
      # For CPU fallback. Will be slower and probably take more memory than the custom CUDA kernel if not executed
      # within a torch.no_grad.
      _, seq_length, _ = key.size()
      output = torch.zeros_like(key)
  
      if state is None:
          num_state = torch.zeros_like(key[:, 0], dtype=torch.float32)
          den_state = torch.zeros_like(key[:, 0], dtype=torch.float32)
          max_state = (
              torch.zeros_like(key[:, 0], dtype=torch.float32) - 1e38
          )
      else:
          num_state, den_state, max_state = state
      # For numerical stability
      #    real_numerator_state = num_state * torch.exp(max_state)
      #    real_denominator_state = den_state * torch.exp(max_state)
  
      time_decay = -torch.exp(time_decay)
  
      for current_index in range(seq_length):
          current_key = key[:, current_index].float()
          current_value = value[:, current_index]
  
          # wkv computation at time t
          max_for_output = torch.maximum(
              max_state, current_key + time_first
          )
          e1 = torch.exp(max_state - max_for_output)
          e2 = torch.exp(current_key + time_first - max_for_output)
          numerator = e1 * num_state + e2 * current_value
          denominator = e1 * den_state + e2
          output[:, current_index] = (numerator / denominator).to(
              output.dtype
          )
  
          # Update state for next iteration
          max_for_state = torch.maximum(
              max_state + time_decay, current_key
          )
          e1 = torch.exp(max_state + time_decay - max_for_state)
          e2 = torch.exp(current_key - max_for_state)
          num_state = e1 * num_state + e2 * current_value
          den_state = e1 * den_state + e2
          max_state = max_for_state
  
      if return_state or state is not None:
          state = [num_state, den_state, max_state]
  
      return output, state
  
imports: import math
  from dataclasses import dataclass
  from pathlib import Path
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import CrossEntropyLoss
  from ...modeling_utils import PreTrainedModel
  from ...utils import ModelOutput, add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, is_bitsandbytes_available, is_ninja_available, is_torch_cuda_available, logging
  from .configuration_rwkv import RwkvConfig
label_desc: Use image node_id 2 for calling a global function with example usage: rwkv_linear_attention_cpu(time_decay, time_first, key, value, state, return_state) and returns: output,  state