code: def test_extract():
      """
      Now we can recover the matching much better by computing the estimated
      encodings for each original image.
      """
      global err, global_matching, guessed_encoded, smatrix
  
      device = torch.device("cuda:1")
  
      print(ooriginal.shape, encoded.shape)
  
      out = []
      for i in range(0, len(ooriginal), 32):
          print(i)
          out.extend(
              inverse(torch.tensor(ooriginal[i : i + 32]).cuda(device))
              .cpu()
              .detach()
              .numpy()
          )
  
      guessed_encoded = np.array(out)
  
      # Now we have to compare each encoded image with every other original image.
      # Do this fast with some matrix multiplies.
  
      out = guessed_encoded.reshape((len(encoded), -1))
      real = encoded[:, global_permutation, :].reshape(
          (len(encoded), -1)
      )
  
      @jax.jit
      def foo(x, y):
          return jn.square(x[:, None] - y[None, :]).sum(2)
  
      smatrix = np.zeros((len(out), len(out)))
  
      B = 500
      for i in range(0, len(out), B):
          print(i)
          for j in range(0, len(out), B):
              smatrix[i : i + B, j : j + B] = foo(
                  out[i : i + B], real[j : j + B]
              )
  
      # And the final time you'l have to look at a min weight matching, I promise.
      row, col = scipy.optimize.linear_sum_assignment(np.array(smatrix))
      r = np.array(smatrix)
  
      print(list(row)[::100])
  
      print("Differences", np.mean(np.argsort(col) != global_matching))
  
      global_matching = np.argsort(col)
  
imports: import collections
  import hashlib
  import time
  import multiprocessing
  import torch
  import numpy
  import torch.nn
  import scipy.stats
  import matplotlib.pyplot
  from PIL import Image
  import jax
  import jax.numpy
  import objax
  import scipy.optimize
  import numpy
  import multiprocessing
label_desc: Use image node_id 17 for calling a global function with example usage: test_extract() and returns: jn