code: def perturb_past(
      past,
      model,
      last,
      unpert_past=None,
      unpert_logits=None,
      accumulated_hidden=None,
      grad_norms=None,
      stepsize=0.01,
      one_hot_bows_vectors=None,
      classifier=None,
      class_label=None,
      loss_type=0,
      num_iterations=3,
      horizon_length=1,
      window_length=0,
      decay=False,
      gamma=1.5,
      kl_scale=0.01,
      device="cuda",
  ):
      # Generate inital perturbed past
      grad_accumulator = [
          (np.zeros(p.shape).astype("float32")) for p in past
      ]
  
      if accumulated_hidden is None:
          accumulated_hidden = 0
  
      if decay:
          decay_mask = torch.arange(
              0.0, 1.0 + SMALL_CONST, 1.0 / (window_length)
          )[1:]
      else:
          decay_mask = 1.0
  
      # TODO fix this comment (SUMANTH)
      # Generate a mask is gradient perturbated is based on a past window
      _, _, _, curr_length, _ = past[0].shape
  
      if curr_length > window_length and window_length > 0:
          ones_key_val_shape = (
              tuple(past[0].shape[:-2])
              + tuple([window_length])
              + tuple(past[0].shape[-1:])
          )
  
          zeros_key_val_shape = (
              tuple(past[0].shape[:-2])
              + tuple([curr_length - window_length])
              + tuple(past[0].shape[-1:])
          )
  
          ones_mask = torch.ones(ones_key_val_shape)
          ones_mask = decay_mask * ones_mask.permute(0, 1, 2, 4, 3)
          ones_mask = ones_mask.permute(0, 1, 2, 4, 3)
  
          window_mask = torch.cat(
              (ones_mask, torch.zeros(zeros_key_val_shape)), dim=-2
          ).to(device)
      else:
          window_mask = torch.ones_like(past[0]).to(device)
  
      # accumulate perturbations for num_iterations
      loss_per_iter = []
      new_accumulated_hidden = None
      for i in range(num_iterations):
          print("Iteration ", i + 1)
          curr_perturbation = [
              torch.from_numpy(p_)
              .requires_grad_(True)
              .to(device=device)
              for p_ in grad_accumulator
          ]
          # make sure p_.grad is not None
          for p_ in curr_perturbation:
              p_.retain_grad()
  
          # Compute hidden using perturbed past
          perturbed_past = list(map(add, past, curr_perturbation))
          _, _, _, curr_length, _ = curr_perturbation[0].shape
          lm_output = model(last, past_key_values=perturbed_past)
          all_logits, all_hidden = (
              lm_output["logits"],
              lm_output["hidden_states"],
          )
          hidden = all_hidden[-1]
          new_accumulated_hidden = (
              accumulated_hidden + torch.sum(hidden, dim=1).detach()
          )
          # TODO: Check the layer-norm consistency of this with trained discriminator (Sumanth)
          logits = all_logits[:, -1, :]
          probs = nn.functional.softmax(logits, dim=-1)
  
          loss = 0.0
          loss_list = []
          if loss_type == PPLM_BOW or loss_type == PPLM_BOW_DISCRIM:
              for one_hot_bow in one_hot_bows_vectors:
                  bow_logits = torch.mm(probs, torch.t(one_hot_bow))
                  bow_loss = -torch.log(torch.sum(bow_logits))
                  loss += bow_loss
                  loss_list.append(bow_loss)
              print(" pplm_bow_loss:", loss.data.cpu().numpy())
  
          if loss_type == 2 or loss_type == 3:
              ce_loss = nn.CrossEntropyLoss()
              # TODO why we need to do this assignment and not just using unpert_past? (Sumanth)
              curr_unpert_past = unpert_past
              curr_probs = torch.unsqueeze(probs, dim=1)
              wte = model.resize_token_embeddings()
              for _ in range(horizon_length):
                  inputs_embeds = torch.matmul(
                      curr_probs, wte.weight.data
                  )
                  lm_output = model(
                      past_key_values=curr_unpert_past,
                      inputs_embeds=inputs_embeds,
                  )
                  curr_all_logits, curr_unpert_past, curr_all_hidden = (
                      lm_output["logits"],
                      lm_output["past_key_values"],
                      lm_output["hidden_states"],
                  )
                  curr_logits = curr_all_logits[:, -1, :]
                  curr_probs = nn.functional.softmax(
                      curr_logits, dim=-1
                  )
                  curr_probs = torch.unsqueeze(curr_probs, dim=1)
                  curr_hidden = curr_all_hidden[-1]
                  new_accumulated_hidden = (
                      new_accumulated_hidden
                      + torch.sum(curr_hidden, dim=1)
                  )
  
              prediction = classifier(
                  new_accumulated_hidden
                  / (curr_length + 1 + horizon_length)
              )
  
              label = torch.tensor(
                  prediction.shape[0] * [class_label],
                  device=device,
                  dtype=torch.long,
              )
              discrim_loss = ce_loss(prediction, label)
              print(
                  " pplm_discrim_loss:", discrim_loss.data.cpu().numpy()
              )
              loss += discrim_loss
              loss_list.append(discrim_loss)
  
          kl_loss = 0.0
          if kl_scale > 0.0:
              unpert_probs = nn.functional.softmax(
                  unpert_logits[:, -1, :], dim=-1
              )
              unpert_probs = (
                  unpert_probs
                  + SMALL_CONST
                  * (unpert_probs <= SMALL_CONST)
                  .float()
                  .to(device)
                  .detach()
              )
              correction = (
                  SMALL_CONST
                  * (probs <= SMALL_CONST).float().to(device).detach()
              )
              corrected_probs = probs + correction.detach()
              kl_loss = kl_scale * (
                  (
                      corrected_probs
                      * (corrected_probs / unpert_probs).log()
                  ).sum()
              )
              print(" kl_loss", kl_loss.data.cpu().numpy())
              loss += kl_loss
  
          loss_per_iter.append(loss.data.cpu().numpy())
          print(" pplm_loss", (loss - kl_loss).data.cpu().numpy())
  
          # compute gradients
          loss.backward()
  
          # calculate gradient norms
          if grad_norms is not None and loss_type == PPLM_BOW:
              grad_norms = [
                  torch.max(
                      grad_norms[index],
                      torch.norm(p_.grad * window_mask),
                  )
                  for index, p_ in enumerate(curr_perturbation)
              ]
          else:
              grad_norms = [
                  (torch.norm(p_.grad * window_mask) + SMALL_CONST)
                  for index, p_ in enumerate(curr_perturbation)
              ]
  
          # normalize gradients
          grad = [
              -stepsize
              * (p_.grad * window_mask / grad_norms[index] ** gamma)
              .data.cpu()
              .numpy()
              for index, p_ in enumerate(curr_perturbation)
          ]
  
          # accumulate gradient
          grad_accumulator = list(map(add, grad, grad_accumulator))
  
          # reset gradients, just to make sure
          for p_ in curr_perturbation:
              p_.grad.data.zero_()
  
          # removing past from the graph
          new_past = []
          for p_ in past:
              new_past.append(p_.detach())
          past = new_past
  
      # apply the accumulated perturbations to the past
      grad_accumulator = [
          torch.from_numpy(p_).requires_grad_(True).to(device=device)
          for p_ in grad_accumulator
      ]
      pert_past = list(map(add, past, grad_accumulator))
  
      return (
          pert_past,
          new_accumulated_hidden,
          grad_norms,
          loss_per_iter,
      )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/examples/research_projects/pplm/run_pplm.py
imports: import argparse
  import json
  from operator import add
  from typing import List, Optional, Tuple, Union
  import numpy
  import torch
  from torch import nn
  from tqdm import trange
  from pplm_classification_head import ClassificationHead
  from transformers import GPT2LMHeadModel, GPT2Tokenizer
  from transformers.file_utils import cached_path
label_desc: Use image node_id 2 for calling a global function with example usage: perturb_past(past, model, last, unpert_past, unpert_logits, accumulated_hidden, grad_norms, stepsize, one_hot_bows_vectors, classifier, class_label, loss_type, num_iterations, horizon_length, window_length, decay, gamma, kl_scale, device) and returns: pert_past,  new_accumulated_hidden,  grad_norms,  loss_per_iter