text: import unittest
  import numpy as np
  import torch
  from torch.nn.modules.distance import CosineSimilarity, PairwiseDistance
  
  from few_shot.utils import *
  from config import PATH
  
  
  class TestDistance(unittest.TestCase):
      def test_query_support_distances(self):
          # Create some dummy data with easily verifiable distances
          q = 1  # 1 query per class
          k = 3  # 3 way classification
          d = 2  # embedding dimension of two
  
          query = torch.zeros([q * k, d], dtype=torch.double)
          query[0] = torch.Tensor([0, 0])
          query[1] = torch.Tensor([0, 1])
          query[2] = torch.Tensor([1, 0])
          support = torch.zeros([k, d], dtype=torch.double)
          support[0] = torch.Tensor([1, 1])
          support[1] = torch.Tensor([1, 2])
          support[2] = torch.Tensor([2, 2])
  
          distances = pairwise_distances(query, support, 'l2')
          self.assertEqual(
              distances.shape, (q * k, k),
              'Output should have shape (q * k, k).'
          )
  
          # Calculate squared distances by iterating through all query-support pairs
          for i, q_ in enumerate(query):
              for j, s_ in enumerate(support):
                  self.assertEqual(
                      (q_ - s_).pow(2).sum(),
                      distances[i, j].item(),
                      'The jth column of the ith row should be the squared distance between the '
                      'ith query sample and the kth query sample'
                  )
  
          # Create some dummy data with easily verifiable distances
          q = 1  # 1 query per class
          k = 3  # 3 way classification
          d = 2  # embedding dimension of two
          query = torch.zeros([q * k, d], dtype=torch.double)
          query[0] = torch.Tensor([1, 0])
          query[1] = torch.Tensor([0, 1])
          query[2] = torch.Tensor([1, 1])
          support = torch.zeros([k, d], dtype=torch.double)
          support[0] = torch.Tensor([1, 1])
          support[1] = torch.Tensor([-1, -1])
          support[2] = torch.Tensor([0, 2])
  
          distances = pairwise_distances(query, support, 'cosine')
  
          # Calculate distances by iterating through all query-support pairs
          for i, q_ in enumerate(query):
              for j, s_ in enumerate(support):
                  self.assertTrue(
                      torch.isclose(1-CosineSimilarity(dim=0)(q_, s_), distances[i, j], atol=2e-8),
                      'The jth column of the ith row should be the squared distance between the '
                      'ith query sample and the kth query sample'
                  )
  
      def test_no_nans_on_zero_vectors(self):
          """Cosine distance calculation involves a divide-through by vector magnitude which
          can divide by zeros to occur.
          """
          # Create some dummy data with easily verifiable distances
          q = 1  # 1 query per class
          k = 3  # 3 way classification
          d = 2  # embedding dimension of two
          query = torch.zeros([q * k, d], dtype=torch.double)
          query[0] = torch.Tensor([0, 0])  # First query sample is all zeros
          query[1] = torch.Tensor([0, 1])
          query[2] = torch.Tensor([1, 1])
          support = torch.zeros([k, d], dtype=torch.double)
          support[0] = torch.Tensor([1, 1])
          support[1] = torch.Tensor([-1, -1])
          support[2] = torch.Tensor([0, 0])  # Third support sample is all zeros
  
          distances = pairwise_distances(query, support, 'cosine')
  
          self.assertTrue(torch.isnan(distances).sum() == 0, 'Cosine distances between 0-vectors should not be nan')
  
  
  class TestAutogradGraphRetrieval(unittest.TestCase):
      def test_retrieval(self):
          """Create a simple autograd graph and check that the output is what is expected"""
          x = torch.ones(2, 2, requires_grad=True)
          y = x + 2
          # The operation on the next line will create two edges because the y variable is
          # y variable is used twice.
          z = y * y
          out = z.mean()
  
          nodes, edges = autograd_graph(out)
  
          # This is quite a brittle test as it will break if the names of the autograd Functions change
          # TODO: Less brittle test
  
          expected_nodes = [
              'MeanBackward1',
              'MulBackward0',
              'AddBackward0',
              'AccumulateGrad',
          ]
  
          self.assertEqual(
              set(expected_nodes),
              set(n.__class__.__name__ for n in nodes),
              'autograd_graph() must return all nodes in the autograd graph.'
          )
  
          # Check for the existence of the expected edges
          expected_edges = [
              ('MulBackward0', 'MeanBackward1'),  # z = y * y, out = z.mean()
              ('AddBackward0', 'MulBackward0'),    # y = x + 2, z = y * y
              ('AccumulateGrad', 'AddBackward0'),   # x = torch.ones(2, 2, requires_grad=True), y = x + 2
          ]
          for e in edges:
              self.assertIn(
                  (e[0].__class__.__name__, e[1].__class__.__name__),
                  expected_edges,
                  'autograd_graph() must return all edges in the autograd graph.'
              )
  
          # Check for two edges between the AddBackward node and the ThMulBackward node
          num_y_squared_edges = 0
          for e in edges:
              if e[0].__class__.__name__ == 'AddBackward0' and e[1].__class__.__name__ == 'MulBackward0':
                  num_y_squared_edges += 1
  
          self.assertEqual(
              num_y_squared_edges,
              2,
              'autograd_graph() must return multiple edges between nodes if they exist.'
          )
  