code: def create_batchwise_out_aux_pairs(
      adj: SparseTensor,
      partitions: List[Union[torch.LongTensor, np.ndarray]],
      prime_indices: Union[torch.LongTensor, np.ndarray],
      topk: int,
      num_outnodeset_per_batch: int = 50,
      alpha: float = 0.2,
      ppr_iterations: int = 50,
  ) -> List[Tuple[np.ndarray, np.ndarray]]:
      def ppr_power_method(
          adj: SparseTensor,
          batch: List[Union[np.ndarray, torch.LongTensor]],
          topk: int,
          num_iter: int,
          alpha: float,
      ) -> List[np.ndarray]:
          topk_neighbors = []
          logits = torch.zeros(
              adj.size(0), len(batch), device=adj.device()
          )  # each column contains a set of output nodes
          for i, tele_set in enumerate(batch):
              logits[tele_set, i] = 1.0 / len(tele_set)
  
          new_logits = logits.clone()
          for i in range(num_iter):
              new_logits = (
                  adj @ new_logits * (1 - alpha) + alpha * logits
              )
  
          inds = new_logits.argsort(0)
          nonzeros = (new_logits > 0).sum(0)
          nonzeros = torch.minimum(
              nonzeros,
              torch.tensor(
                  [topk], dtype=torch.int64, device=adj.device()
              ),
          )
          for i in range(new_logits.shape[1]):
              topk_neighbors.append(
                  inds[-nonzeros[i] :, i].cpu().numpy()
              )
  
          return topk_neighbors
  
      device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
      if isinstance(prime_indices, Tensor):
          prime_indices = prime_indices.cpu().numpy()
  
      adj = adj.to(device)
  
      cur_output_nodes = []
      loader = []
  
      pbar = tqdm(range(len(partitions)))
      pbar.set_description("Processing topic-sensitive PPR batches")
      for n in pbar:
          part = partitions[n]
          if isinstance(part, Tensor):
              part = part.cpu().numpy()
  
          primes_in_part, *_ = np.intersect1d(
              part,
              prime_indices,
              assume_unique=True,
              return_indices=True,
          )
          if len(primes_in_part):  # no output nodes in this partition
              cur_output_nodes.append(primes_in_part)
  
          # accumulate enough output nodes to make good use of GPU memory
          if (
              len(cur_output_nodes) >= num_outnodeset_per_batch
              or n == len(partitions) - 1
          ):
              topk_neighbors = ppr_power_method(
                  adj, cur_output_nodes, topk, ppr_iterations, alpha
              )
              for i in range(len(cur_output_nodes)):
                  # force output nodes to be aux nodes
                  auxiliary_nodes = np.union1d(
                      cur_output_nodes[i], topk_neighbors[i]
                  )
                  loader.append((cur_output_nodes[i], auxiliary_nodes))
              cur_output_nodes = []
  
          if torch.cuda.is_available():
              torch.cuda.empty_cache()
  
      return loader
  
imports: import logging
  import math
  from typing import Callable, Iterator, List, NamedTuple, Optional, Tuple, Union
  import numpy
  import scipy.sparse
  import torch
  from torch import Tensor
  from tqdm import tqdm
  from torch_geometric.data import Data
  from torch_geometric.typing import SparseTensor
  from torch_geometric.utils import get_ppr, is_undirected, subgraph
label_desc: Use image node_id 6 for calling a global function with example usage: create_batchwise_out_aux_pairs(adj, partitions, prime_indices, topk, num_outnodeset_per_batch, alpha, ppr_iterations) and returns: loader, topk_neighbors