YaRN Interpolation Ramp

Problem

Compute the YaRN interpolation ramp for one RoPE pair index: indices at or below alpha keep their frequency, indices at or above beta are fully interpolated, and values in between ramp linearly.

The signature is yarn_ramp(pair_index, alpha, beta). Return a float in [0.0, 1.0]; compute (pair_index - alpha) / (beta - alpha) for the middle range.

Hint

Dimensions that rotate quickly need no interpolation, while slow long-wavelength dimensions need full interpolation.

Test cases

case 1

in: [0,0,32]

expected: 0

case 2

in: [16,0,32]

expected: 0.5

case 3

in: [32,0,32]

expected: 1

case 4

in: [8,4,12]

expected: 0.5

Starter codePython

def yarn_ramp(pair_index, alpha, beta):
    # Your code here
    pass

4 test cases · run your code against them in the editor