//  [python]
def anderson(f, x0, X, F, H, y, args, m=5, lam=1e-3, max_iter=15, tol=1e-2, beta = 1.0, logger=None):
    """ Anderson acceleration for fixed point iteration. """
    with torch.no_grad():
        bsz, ch, h0, w0 = x0.shape
        
        t1 = time.time()

        X[:,0] = x0.view(bsz, -1)
        F[:,0], _ = f(xt=x0.view(x0.shape), **args)

        X[:,1] = F[:,0].view(bsz, -1)
        F[:,1], _ = f(xt=F[:,0].view(x0.shape), **args)

        H[:,0,1:] = H[:,1:,0] = 1
        y[:,0] = 1
        
        t2 = time.time()
        print("Intial set up", t2-t1)
        time_logger = {
            "setup": t2 - t1,
            "bmm": 0,
            "solve": 0,
            "forward call-unet": 0,
            "total_time_per_iter": 0
        }

        iter_count = 0
        log_metrics = {}
        res = []
        norm_res = []
        for k in range(2, max_iter):
            n_ = min(k, m)
            G = F[:,:n_]-X[:,:n_]
            
            t3 = time.time()
            H[:,1:n_+1,1:n_+1] = torch.bmm(G,G.transpose(1,2)) + lam*torch.eye(n_, dtype=x0.dtype,device=x0.device)[None]
            t4 = time.time()

            alpha = torch.solve(y[:,:n_+1], H[:,:n_+1,:n_+1])[0][:, 1:n_+1, 0]   # (bsz x n)
            t5 = time.time()
            X[:,k%m] = beta * (alpha[:,None] @ F[:,:n_])[:,0] + (1-beta)*(alpha[:,None] @ X[:,:n_])[:,0]
            F[:,k%m], log_metrics = f(xt=X[:,k%m].view(x0.shape), **args)
            t6 = time.time()

            residual = (F[:,k%m] - X[:,k%m]).norm().item()
            normalized_residual = (F[:,k%m] - X[:,k%m]).norm().item()/(1e-5 + F[:,k%m].norm()).item()

            t7 = time.time()
            res.append(residual)
            norm_res.append(normalized_residual)
            
            time_logger["bmm"] += (t4 - t3)
            time_logger["solve"] += (t5 - t4)
            time_logger["forward call-unet"] += (t6 - t5)
            time_logger["total_time_per_iter"] += (t7 - t3)
            iter_count += 1

            ### TODO: break out early for norm_res
            if (norm_res[-1] < tol):
                print("Breaking out early at {}".format(k))
                break

            #print("{}/{} Residual {} tol {} ".format(k, max_iter, res[-1], tol))
            if logger is not None:
                log_metrics["residual"] = residual
                log_metrics["normalized_residual"] = normalized_residual

                log_metrics["alpha"] = torch.norm(alpha, dim=-1).mean()
                log_metrics["samples"] = [wandb.Image(X[:, k%m].view_as(x0).to('cpu')[ts]) for ts in args['plot_timesteps']]
                log_metrics["setup"] = time_logger['setup']
                log_metrics["total_time_per_iter"] = time_logger['total_time_per_iter'] / iter_count
                log_metrics["total_time"] = t7 - t1
                for key, val in time_logger.items():
                    if key not in log_metrics: 
                        log_metrics[f"avg-{key}"] = val / iter_count
                        log_metrics[key] = val
                log_metrics["perc_time_forward_call"] = time_logger["forward call-unet"] * 100 / time_logger["total_time_per_iter"]
                logger(log_metrics)
    x_eq = X[:,k%m].view_as(x0)[args['gather_idx']].to('cpu')
    return x_eq