code: def cosine_similarity(test, train, replace_nan=0) -> Tensor:
      r"""
      Calculates the pairwise cosine similarity for batches of feature vectors.
      Tensors test and train have shape (batch_size_1, *), and (batch_size_2, *).
      Returns pairwise cosine similarity Tensor of shape (batch_size_1, batch_size_2).
      """
      test = test.view(test.shape[0], -1)
      train = train.view(train.shape[0], -1)
  
      if common._parse_version(torch.__version__) <= (1, 6, 0):
          test_norm = torch.norm(test, p=None, dim=1, keepdim=True)
          train_norm = torch.norm(train, p=None, dim=1, keepdim=True)
      else:
          test_norm = torch.linalg.norm(
              test, ord=2, dim=1, keepdim=True
          )
          train_norm = torch.linalg.norm(
              train, ord=2, dim=1, keepdim=True
          )
  
      test = torch.where(
          test_norm != 0.0, test / test_norm, Tensor([replace_nan])
      )
      train = torch.where(
          train_norm != 0.0, train / train_norm, Tensor([replace_nan])
      ).T
  
      similarity = torch.mm(test, train)
      return similarity
  
file_path: captum/captum/influence/_core/similarity_influence.py
imports: import warnings
  from functools import partial
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  import captum._utils.common
  import torch
  from captum._utils.av import AV
  from captum.attr import LayerActivation
  from captum.influence._core.influence import DataInfluence
  from torch import Tensor
  from torch.nn import Module
  from torch.utils.data import DataLoader, Dataset
label_desc: Use image node_id 2 for calling a global function with example usage: cosine_similarity(test, train, replace_nan) and returns: similarity