//  [python]
def butterfly_projection_cov(teacher_module, input_cov, butterfly_structure='odo_1',
                             n_Adam_steps=20000, n_LBFGS_steps=50):
    teacher_module = teacher_module.float()
    input_cov = input_cov.float()
    try:
        in_channels = teacher_module.in_channels
        out_channels = teacher_module.out_channels
    except:
        raise ValueError("Only convolutional layers currently supported.")
    param = butterfly_structure.split('_')[0]
    nblocks = 0 if len(butterfly_structure.split('_')) <= 1 else int(butterfly_structure.split('_')[1])
    residual = False if len(butterfly_structure.split('_')) <= 2 else (butterfly_structure.split('_')[2] == 'res')
    student_module = Butterfly1x1Conv(in_channels, out_channels,
        bias=False, tied_weight=False, ortho_init=True,
        param=param, nblocks=nblocks)
    student_module = student_module.to(input_cov.device)

    with torch.no_grad():
        try:  # torch.linalg.eigh sometimes fail to converge on CUDA
            Sigma, U = torch.linalg.eigh(input_cov)
        except:  # Move to CPU and use numpy's function
            Sigma, U = np.linalg.eigh(input_cov.cpu().numpy())
            Sigma = torch.tensor(Sigma, dtype=input_cov.dtype, device=input_cov.device)
            U = torch.tensor(U, dtype=input_cov.dtype, device=input_cov.device)
        Sigma = Sigma.clamp(0)  # avoid small negative eigenvalues
        input = torch.diag(Sigma.sqrt()) @ U.t()
        input = input.reshape(in_channels, in_channels, 1, 1)  # To be compatible with conv2d
        target = teacher_module(input)
        # Normalize input so that output has MSE 1.0
        input /= (target ** 2).mean().sqrt()
        target = teacher_module(input)

    def loss_fn():
        output = student_module(input)
        if residual:
            if output.shape[1] == 2 * input.shape[1]:
                b, c, h, w = input.shape
                output = (output.reshape(b, 2, c, h, w) + input.reshape(b, 1, c, h, w)).reshape(b, 2 * c, h, w)
            else:
                output = output + input
        return F.mse_loss(output, target)

    optimizer = optim.Adam(student_module.parameters())
    student_module.train()
    for _ in range(n_Adam_steps):
        optimizer.zero_grad()
        loss = loss_fn()
        loss.backward()
        optimizer.step()

    optimizer = optim.LBFGS(filter(lambda p: p.requires_grad, student_module.parameters()),
                            tolerance_grad=1e-7,  # Pytorch 1.2 sets this too high https://github.com/pytorch/pytorch/pull/25240
                            line_search_fn='strong_wolfe')
    def closure():
        optimizer.zero_grad()
        loss = loss_fn()
        loss.backward()
        return loss
    for i in range(n_LBFGS_steps):
        loss = optimizer.step(closure)

    if torch.distributed.is_initialized():
        # Get the model from the process with the lowest loss.
        # Losses could be different due to different initialization of student_module.
        all_losses = [torch.empty_like(loss) for _ in range(torch.distributed.get_world_size())]
        torch.distributed.all_gather(all_losses, loss)
        best_rank = min(range(len(all_losses)), key=lambda i: all_losses[i])
        loss = all_losses[best_rank]
        for p in student_module.parameters():
            torch.distributed.broadcast(p, best_rank)

    return student_module, loss.item()