//  [python]
def decompose_weights(model, low_rank_model, rank_list, rank_ratio, args):
    # SVD version
    reconstructed_aggregator = []
    
    if args.arch == "vgg19":
        layer_counter = 0
        if args.mode == "pufferfish":
            fullrank_layer_range = 54
        else:
            fullrank_layer_range = 24

        for item_index, (param_name, param) in enumerate(model.state_dict().items()):
            if len(param.size()) == 4 and item_index not in range(0, fullrank_layer_range):
                # resize --> svd --> two layer
                param_reshaped = param.view(param.size()[0], -1)
                rank = min(param_reshaped.size()[0], param_reshaped.size()[1])
                u, s, v = torch.linalg.svd(param_reshaped)

                if args.mode == "lowrank":
                    sliced_rank = rank_list[layer_counter]
                elif args.mode in ("baseline", "pufferfish"):
                    sliced_rank = int(rank/rank_ratio)
                else:
                    raise NotImplementedError("Unsupported mode ...")

                #u_weight = u * torch.sqrt(s) # alternative implementation: u_weight_alt = torch.mm(u, torch.diag(torch.sqrt(s)))
                #v_weight = torch.sqrt(s) * v # alternative implementation: v_weight_alt = torch.mm(torch.diag(torch.sqrt(s)), v.t())
                # alternative implementation
                u_weight = torch.matmul(u, torch.diag(torch.sqrt(s)))
                v_weight = torch.matmul(torch.diag(torch.sqrt(s)), v.t()).t()
                
                #print("layer indeix: {}, dist u u_alt:{}, dist v v_alt: {}".format(item_index, torch.dist(u_weight, u_weight_alt), torch.dist(v_weight.t(), v_weight_alt)))
                #print("layer indeix: {}, dist u u_alt:{}, dist v v_alt: {}".format(item_index, torch.equal(u_weight, u_weight_alt), torch.equal(v_weight.t(), v_weight_alt)))
                u_weight_sliced, v_weight_sliced = u_weight[:, 0:sliced_rank], v_weight[:, 0:sliced_rank]

                u_weight_sliced_shape, v_weight_sliced_shape = u_weight_sliced.size(), v_weight_sliced.size()

                model_weight_v = u_weight_sliced.view(u_weight_sliced_shape[0],
                                                      u_weight_sliced_shape[1], 1, 1)
                
                model_weight_u = v_weight_sliced.t().view(v_weight_sliced_shape[1], 
                                                          param.size()[1], 
                                                          param.size()[2], 
                                                          param.size()[3])

                reconstructed_aggregator.append(model_weight_u)
                reconstructed_aggregator.append(model_weight_v)
                layer_counter += 1
            else:
                reconstructed_aggregator.append(param)
                
                
        model_counter = 0
        reload_state_dict = {}
        for item_index, (param_name, param) in enumerate(low_rank_model.state_dict().items()):
            # print("#### {}, {}, recons agg: {}, param: {}".format(item_index, param_name, 
            #                                                                         reconstructed_aggregator[model_counter].size(),
            #                                                                        param.size()))
            if "batch_norm" in param_name and "_u" in param_name:
                reload_state_dict[param_name] = param
            else:
                assert (reconstructed_aggregator[model_counter].size() == param.size())
                reload_state_dict[param_name] = reconstructed_aggregator[model_counter]
                model_counter += 1

    elif args.arch == "resnet18":
        layer_counter = 0

        if args.mode == "pufferfish":
            fullrank_layer_range = 18
        elif args.mode == "baseline":
            fullrank_layer_range = RESNET18_FR_BLOCKS_IDX_MAP[0]
        elif args.mode == "lowrank":
            fullrank_layer_range = 25
        else:
            raise NotImplementedError("Unsupported training mode ...")

        for item_index, (param_name, param) in enumerate(model.state_dict().items()):
            #if len(param.size()) == 4 and item_index not in range(0, 18) and ".shortcut." not in param_name:
            #if len(param.size()) == 4 and item_index not in range(0, 13) and ".shortcut." not in param_name:
            #if len(param.size()) == 4 and item_index not in range(0, 25) and ".shortcut." not in param_name: # two full-rank blocks
            #if len(param.size()) == 4 and item_index not in range(0, 48) and ".shortcut." not in param_name: # three full-rank blocks
            if len(param.size()) == 4 and item_index not in range(0, fullrank_layer_range) and ".shortcut." not in param_name: # three full-rank blocks
            #if len(param.size()) == 4 and item_index not in range(0, 60) and ".shortcut." not in param_name:
                # resize --> svd --> two layer
                param_reshaped = param.view(param.size()[0], -1)
                rank = min(param_reshaped.size()[0], param_reshaped.size()[1])
                u, s, v = torch.linalg.svd(param_reshaped)

                if args.mode == "lowrank":
                    sliced_rank = rank_list[layer_counter]
                elif args.mode in ("baseline", "pufferfish"):
                    sliced_rank = int(rank/rank_ratio)
                else:
                    raise NotImplementedError("Unsupported mode ...")
                    
                #u_weight = u * torch.sqrt(s)
                #v_weight = torch.sqrt(s) * v
                u_weight = torch.matmul(u, torch.diag(torch.sqrt(s)))
                #v_weight = torch.mm(torch.diag(torch.sqrt(s)), v.t()).t()
                v_weight = torch.matmul(torch.diag(torch.sqrt(s)), v.t()).t()

                u_weight_sliced, v_weight_sliced = u_weight[:, 0:sliced_rank], v_weight[:, 0:sliced_rank]

                u_weight_sliced_shape, v_weight_sliced_shape = u_weight_sliced.size(), v_weight_sliced.size()

                model_weight_v = u_weight_sliced.view(u_weight_sliced_shape[0],
                                                      u_weight_sliced_shape[1], 1, 1)
                
                model_weight_u = v_weight_sliced.t().view(v_weight_sliced_shape[1], 
                                                          param.size()[1], 
                                                          param.size()[2], 
                                                          param.size()[3])

                reconstructed_aggregator.append(model_weight_u)
                reconstructed_aggregator.append(model_weight_v)
                layer_counter += 1
            else:
                reconstructed_aggregator.append(param)
                
        model_counter = 0
        reload_state_dict = {}
        for item_index, (param_name, param) in enumerate(low_rank_model.state_dict().items()):
            #print("#### {}, {}, recons agg: {}， param: {}".format(item_index, param_name, 
            #                                                                        reconstructed_aggregator[model_counter].size(),
            #                                                                       param.size()))
            if ".bn1_u." in param_name or ".bn2_u." in param_name:
                reload_state_dict[param_name] = param
            else:
                assert (reconstructed_aggregator[model_counter].size() == param.size())
                reload_state_dict[param_name] = reconstructed_aggregator[model_counter]
                model_counter += 1
    else:
        raise NotImplementedError("Unsupported model arch ...")
    
    low_rank_model.load_state_dict(reload_state_dict)
    return low_rank_model