Method extra_repr

k_diffusion/models/axial_rope.py:95–97

(self)

        self.freqs_w = nn.Parameter(log_freqs.clone())

    def extra_repr(self):
        dim = (self.freqs_h.shape[-1] + self.freqs_w.shape[-1]) * 2
        return f"dim={dim}, n_heads={self.n_heads}, start_index={self.start_index}"

    def get_freqs(self, pos):
        if pos.shape[-1] != 2: