code: def exact_diffusion(
      X,
      y,
      w_opt,
      loss,
      maxite=2000,
      alpha=1e-1,
      use_Abar=True,
      **kwargs,
  ):
      if loss == "logistic_regression":
          rho = kwargs.get("rho", 1e-1)
      elif loss == "linear_regression":
          rho = 0
      else:
          raise NotImplementedError(
              "Task not supported. This example only supports"
              + " linear_regression and logistic_regression"
          )
  
      topology = bf.load_topology()
      self_weight, neighbor_weights = topology_util.GetRecvWeights(
          topology, bf.rank()
      )
  
      if bf.rank() == 0:
          print("self weights with A: {}\n".format(self_weight))
          print("neighbor weights with A:\n")
          for k, v in neighbor_weights.items():
              print(k, v)
  
      w = torch.zeros(n, 1, dtype=torch.double, requires_grad=True)
      phi, psi, psi_prev = w.clone(), w.clone(), w.clone()
      mse = []
  
      # construct A_bar
      if use_Abar:
          self_weight = (self_weight + 1) / 2
          for k, v in neighbor_weights.items():
              neighbor_weights[k] = v / 2
  
      for i in range(maxite):
          # calculate loccal gradient via pytorch autograd
          loss_step(
              X,
              y,
              w,
              tensor_name="neighbor.allreduce.local_variable",
              loss=loss,
              rho=rho,
          )
  
          # exact diffusion
          with torch.no_grad():
              psi = w - alpha * w.grad.data
              phi = psi + w.data - psi_prev
              w.data = bf.neighbor_allreduce(
                  phi,
                  self_weight=self_weight,
                  src_weights=neighbor_weights,
                  name="local variable",
              )
              psi_prev = psi.clone()
              w.grad.data.zero_()
  
              # record convergence
              if bf.rank() == 0:
                  mse.append(torch.norm(w.data - w_opt.data, p=2))
  
      return w, mse
  
file_path: bluefog/examples/pytorch_optimization.py
imports: import os
  import torch
  import matplotlib
  import matplotlib.pyplot
  import argparse
  import bluefog.torch
  from bluefog.common import topology_util
  import networkx
label_desc: Use image node_id 7 for calling a global function with example usage: exact_diffusion(X, y, w_opt, loss, maxite, alpha, use_Abar) and returns: w,  mse