Method forward

k_diffusion/models/axial_rope.py:107–109

(self, x, pos)

        return freqs.transpose(-2, -3)

    def forward(self, x, pos):
        freqs = self.get_freqs(pos)
        return apply_rotary_emb(freqs, x, self.start_index)