(self, dim, n_heads, start_index=0, freqs_init=freqs_pixel_log(max_freq=10.0))
| 85 | |
| 86 | class AxialRoPE(nn.Module): |
| 87 | def __init__(self, dim, n_heads, start_index=0, freqs_init=freqs_pixel_log(max_freq=10.0)): |
| 88 | super().__init__() |
| 89 | self.n_heads = n_heads |
| 90 | self.start_index = start_index |
| 91 | log_freqs = freqs_init((n_heads, dim // 4)) |
| 92 | self.freqs_h = nn.Parameter(log_freqs.clone()) |
| 93 | self.freqs_w = nn.Parameter(log_freqs.clone()) |
| 94 | |
| 95 | def extra_repr(self): |
| 96 | dim = (self.freqs_h.shape[-1] + self.freqs_w.shape[-1]) * 2 |
nothing calls this directly
no test coverage detected