code: def forward(self, query_layer: torch.Tensor, key_layer: torch.Tensor):
      return self.rope_fn(
          self.cos_cached, self.sin_cached, query_layer, key_layer
      )
  
imports: import torch
  import math
  from typing import Tuple
  from einops import repeat
label_desc: Use image node_id 4 for calling the Rotary obj's underlying member method code with example usage: obj.forward(query_layer, key_layer) and returns: self