import torch
from typing import Dict, List


def merge_lora_weight(tensors_A, tensors_B):
    lora_A = torch.concat(tensors_A, dim=0)
    lora_B = torch.concat(tensors_B, dim=1)
    return lora_A, lora_B


def merge_lora(loras: List[Dict[str, torch.Tensor]], alpha=1):
    lora_merged = {}
    keys = [i for i in loras[0].keys() if ".lora_A." in i]
    for key in keys:
        tensors_A = [lora[key] for lora in loras]
        tensors_B = [lora[key.replace(".lora_A.", ".lora_B.")] for lora in loras]
        lora_A, lora_B = merge_lora_weight(tensors_A, tensors_B)
        lora_merged[key] = lora_A * alpha
        lora_merged[key.replace(".lora_A.", ".lora_B.")] = lora_B
    return lora_merged