text: """
  Semi-Amortized VAE optimization procedure.
  
  Adapted from: https://github.com/harvardnlp/sa-vae
  which provides code for the paper "Semi-Amortized Variational Autoencoders"
  Yoon Kim, Sam Wiseman, Andrew Miller, David Sontag, Alexander Rush.
  
  """
  
  import torch
  import numpy as np
  
  
  class OptimN2N:
      def __init__(self, loss_fn, model, model_update_params,
                   beta,
                   lr=[1],
                   iters=20,
                   acc_param_grads=True,
                   max_grad_norm=0,
                   eps=0.00001,
                   momentum=0.5):
  
          self.iters = iters
          self.lr = lr
          self.beta = beta
          self.loss_fn = loss_fn
          self.eps = eps
          self.max_grad_norm = max_grad_norm
          self.model = model
          self.momentum = momentum
          self.acc_param_grads = acc_param_grads
          if self.acc_param_grads:
              self.params = model_update_params
              self.param_grads = [torch.zeros([self.iters] + list(p.size())).type_as(p.data)
                                  for p in self.params]
  
      def forward(self, input, y, verbose=False):
          self.seeds = np.random.randint(3435, size=self.iters)
          return self.forward_mom(input, y, verbose)
  
      def backward(self, grad_output, verbose=False):
          grads = self.backward_mom(grad_output, verbose)
          return grads
  
      def grad_norm(self, g_list, g_norm_list):
          for g, g_norm in zip(g_list, g_norm_list):
              g_norm2 = (g ** 2).sum(1) ** 0.5
              g.div_(g_norm2.unsqueeze(1).expand_as(g)).mul_(g_norm.unsqueeze(1).expand_as(g))
  
      def clip_grad_norm(self, parameters, max_norm, norm_type=2):
          if len(parameters) > 0:
              max_norm = float(max_norm)
              norm_type = float(norm_type)
              if norm_type == float('inf'):
                  total_norm = max(p.abs().max() for p in parameters)
              else:
                  total_norm = 0
                  for p in parameters:
                      param_norm = p.norm(norm_type)
                      total_norm += param_norm ** norm_type
                  total_norm = total_norm ** (1. / norm_type)
                  clip_coef = max_norm / (total_norm + 1e-6)
                  if clip_coef < 1:
                      for p in parameters:
                          p.mul_(clip_coef)
  
      def forward_mom(self, input, y, verbose=False):
          self.y = y
          self.input_grads = [torch.zeros([self.iters] + list(x.size())).type_as(x.data) for x in input]
          self.mom_params = [torch.zeros(x.size()).type_as(x) for x in self.input_grads]
          self.input_cache = [torch.zeros(x.size()).type_as(x) for x in self.input_grads]
          self.all_z = []
          if self.acc_param_grads:
              for p in self.param_grads:
                  p.zero_()
          for k in range(self.iters):
              self.all_z.append(torch.randn_like(input[0]))
              torch.manual_seed(int(self.seeds[k]))
              loss = self.loss_fn(input, self.y, self.model, self.beta, self.all_z[k])
  
              if self.acc_param_grads:
                  all_input_params = input + self.params
              else:
                  all_input_params = input
              all_grads_k = torch.autograd.grad(loss, all_input_params, retain_graph=True)
              input_grad_k = all_grads_k[:len(input)]
              param_grads_k = all_grads_k[len(input):]
  
              if self.max_grad_norm > 0:
                  self.clip_grad_norm([input_grad_k[0].data], self.max_grad_norm)
                  self.clip_grad_norm([input_grad_k[1].data], self.max_grad_norm)
  
              if self.acc_param_grads:
                  for i, p in enumerate(param_grads_k):
                      self.param_grads[i][k].copy_(p.data)
              for i, x_grad_k in enumerate(input_grad_k):
                  self.input_cache[i][k].copy_(input[i].data)
                  self.input_grads[i][k].copy_(x_grad_k.data)
  
              for i in range(len(self.mom_params)):
                  if k == 0:
                      self.mom_params[i][k] = -input_grad_k[i].data
                  else:
                      self.mom_params[i][k] = self.mom_params[i][k - 1] * self.momentum - input_grad_k[i].data
              lr_k_list = [lr for lr in self.lr]
              input = [(x.data + lr_k * p[k]).clone().detach().requires_grad_(True) for x, p, lr_k in
                       zip(input, self.mom_params, lr_k_list)]
              if verbose:
                  print('mom', k, loss.item())
          return input
  
      def backward_mom(self, grad_output, verbose=False):
          input_kp1_grad = [g.data for g in grad_output]
          p_kp1_grad = [torch.zeros(x.size()).type_as(x) for x in input_kp1_grad]
          rev_iters = self.iters
          for k in reversed(range(rev_iters)):
              lr_k_list = [lr for lr in self.lr]
              input_k_grad = input_kp1_grad
              p_kp1_grad = [p + lr_k * x for p, x, lr_k in zip(p_kp1_grad, input_kp1_grad, lr_k_list)]
              input_k_rv = []
              input_H_xx_v = []
              r = self.eps
              for i in range(len(p_kp1_grad)):
                  v = p_kp1_grad[i]
                  x_k = self.input_cache[i][k]
                  x_k_rv = (x_k + r * v).type_as(x_k).clone().detach().requires_grad_(True)
                  input_k_rv.append(x_k_rv)
              if self.acc_param_grads:
                  all_input_params = input_k_rv + self.params
              else:
                  all_input_params = input_k_rv
              torch.manual_seed(int(self.seeds[k]))
              loss = self.loss_fn(input_k_rv, self.y, self.model, self.beta, self.all_z[k])
              all_grads_rv_k = torch.autograd.grad(loss, all_input_params, retain_graph=True)
  
              if self.max_grad_norm > 0:
                  self.clip_grad_norm([g.data for g in all_grads_rv_k], self.max_grad_norm)
  
              input_grads_rv_k = all_grads_rv_k[:len(input_k_rv)]
              param_grads_rv_k = all_grads_rv_k[len(input_k_rv):]
  
              if self.acc_param_grads:
                  H_wx_v_list = []
                  for i, p_grad_rv_k in enumerate(param_grads_rv_k):
                      H_wx_v = (p_grad_rv_k.data - self.param_grads[i][k]) / r
                      H_wx_v_list.append(H_wx_v)
                      if self.params[i].grad is None:
                          self.params[i].grad = torch.zeros(self.params[i].size()).type_as(
                              self.params[i].data)
                  if self.max_grad_norm > 0:
                      self.clip_grad_norm(H_wx_v_list, self.max_grad_norm)
                  for i in range(len(self.params)):
                      self.params[i].grad.data += -H_wx_v_list[i]
              for i, x_k_rv_grad in enumerate(input_grads_rv_k):
                  H_xx_v = (x_k_rv_grad.data - self.input_grads[i][k]) / r
                  input_H_xx_v.append(H_xx_v)
              input_kp1_grad = [x_kp1_grad - H_xx_v
                                for (x_kp1_grad, H_xx_v) in zip(input_kp1_grad, input_H_xx_v)]
              if self.max_grad_norm > 0:
                  self.clip_grad_norm(input_kp1_grad, self.max_grad_norm)
              p_kp1_grad = [p.mul_(self.momentum) for p in p_kp1_grad]
              if verbose:
                  print('mom', k, input_kp1_grad[0][0].norm())
          return input_kp1_grad
  
  