Function init

k_diffusion/models/axial_rope.py:72–74

(shape)

def freqs_pixel(max_freq=10.0):
    def init(shape):
        freqs = torch.linspace(1.0, max_freq / 2, shape[-1]) * math.pi
        return freqs.log().expand(shape)
    return init