LoRA Merge Weights

Problem

Merge a LoRA adapter into its base weight: W_merged = W + (alpha / rank) * (B @ A), where W is (d_out, d_in), A is (rank, d_in) and B is (d_out, rank).

The signature is lora_merge(W, A, B, alpha, rank).

Hint

Merging removes inference overhead because the adapter becomes part of the base matrix.

Test cases

case 1

in: [[[1,0],[0,1]],[[1,0]],[[1],[0]],2,1]

expected: [[3,0],[0,1]]

case 2

in: [[[0,0],[0,0]],[[1,0],[0,1]],[[1,0],[0,1]],1,2]

expected: [[0.5,0],[0,0.5]]

case 3

in: [[[1,2],[3,4]],[[1,1]],[[1],[1]],4,2]

expected: [[3,4],[5,6]]

case 4

in: [[[5]],[[2]],[[3]],1,3]

expected: [[7]]

Starter codePython

def lora_merge(W, A, B, alpha, rank):
    # Your code here
    pass

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