code: def __init__(
      self,
      layer_idx: int,
      num_layers: int,
      num_tokens: int,
      in_dim: int,
      latent_dim: int,
  ):
      super().__init__(in_dim, num_tokens)
      latent_dim = latent_dim or in_dim
      self.ci = latent_dim
  
      self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
      self.key = nn.Linear(in_dim, latent_dim, bias=False)
      self.value = nn.Linear(in_dim, latent_dim, bias=False)
      self.receptance = nn.Linear(in_dim, latent_dim, bias=False)
      self.output = nn.Linear(latent_dim, in_dim, bias=False)
  
      with torch.no_grad():
          for k, v in self.state_dict().items():
              for zero_k in ["key", "receptance", "output"]:
                  if zero_k in k:
                      nn.init.zeros_(v)
                      break
              else:
                  nn.init.orthogonal_(v)
  
          _0_to_1 = layer_idx / (num_layers - 1)
          _1_to_almost0 = 1.0 - (layer_idx / num_layers)
  
          f1_begin = 3.0
          f1_end = 1.2
          f2_begin = 0.65
          f2_end = 0.4
          decay = torch.ones(latent_dim, 1)
          first_sa_layer_id = int(num_layers > 2)
          for h in range(latent_dim):
              mul = layer_idx - first_sa_layer_id
              divider = num_layers - 1 - first_sa_layer_id
              f1 = f1_begin + mul / divider * (f1_end - f1_begin)
              f2 = f2_begin + mul / divider * (f2_end - f2_begin)
              if layer_idx == first_sa_layer_id:
                  f1 += 0.5
              elif layer_idx == num_layers - 2:
                  f2 = 0.4
              elif layer_idx == num_layers - 1:
                  f2 = 0.37
              decay[h][0] = (
                  math.pow(f2, h / (latent_dim - 1) * 7.0) * f1
              )
          # Cl, 1
          self.info_decay = nn.Parameter(decay.log())
          # 1, T
          time_curve = torch.arange(-num_tokens + 1, 1)[None]
          self.register_buffer("time_curve", time_curve)
  
          x = torch.ones(1, 1, in_dim)
          for i in range(in_dim):
              x[0, 0, i] = i / in_dim
          self.time_mix_k = nn.Parameter(torch.pow(x, _1_to_almost0))
          self.time_mix_v = nn.Parameter(
              torch.pow(x, _1_to_almost0) + 0.3 * _0_to_1
          )
          self.time_mix_r = nn.Parameter(
              torch.pow(x, 0.5 * _1_to_almost0)
          )
  
imports: import math
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from typing import Any
  from typing import Tuple
  from typing import Optional
  from torch.fft import fft
  from cftool.misc import safe_execute
  from .schema import TokenMixerBase
  from .channel_mixers import FeedForward
  from ..common import Lambda
  from ..attentions import Attention
label_desc: Use image node_id 1 to create a new RWKVTokenMixer object from inherited base classes: TokenMixerBase with example: obj = RWKVTokenMixer(layer_idx, num_layers, num_tokens, in_dim, latent_dim)