code: def __init__(self, dim, max_freq=10):
      super().__init__()
      self.dim = dim
      scales = torch.linspace(1.0, max_freq / 2, self.dim // 4)
      self.register_buffer("scales", scales)
  
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 1 to create a new AxialRotaryEmbedding object from inherited base classes: nn with example: obj = AxialRotaryEmbedding(dim, max_freq)