code: def forward(self, x):
      device, dtype, n = x.device, x.dtype, int(sqrt(x.shape[-2]))
  
      seq = torch.linspace(-1.0, 1.0, steps=n, device=device)
      seq = seq.unsqueeze(-1)
  
      scales = self.scales[
          (*((None,) * (len(seq.shape) - 1)), Ellipsis)
      ]
      scales = scales.to(x)
  
      seq = seq * scales * pi
  
      x_sinu = repeat(seq, "i d -> i j d", j=n)
      y_sinu = repeat(seq, "j d -> i j d", i=n)
  
      sin = torch.cat((x_sinu.sin(), y_sinu.sin()), dim=-1)
      cos = torch.cat((x_sinu.cos(), y_sinu.cos()), dim=-1)
  
      sin, cos = map(
          lambda t: rearrange(t, "i j d -> (i j) d"), (sin, cos)
      )
      sin, cos = map(
          lambda t: repeat(t, "n d -> () n (d j)", j=2), (sin, cos)
      )
      return sin, cos
  
imports: from math import sqrt, pi, log
  import torch
  from torch import nn, einsum
  import torch.nn.functional
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 2 for calling the AxialRotaryEmbedding obj's underlying member method code with example usage: obj.forward(x) and returns: sin,  cos