text: import torch
  
  from ..utils.base_model import BaseModel
  
  
  def find_nn(sim, ratio_thresh, distance_thresh):
      sim_nn, ind_nn = sim.topk(2 if ratio_thresh else 1, dim=-1, largest=True)
      dist_nn = 2 * (1 - sim_nn)
      mask = torch.ones(ind_nn.shape[:-1], dtype=torch.bool, device=sim.device)
      if ratio_thresh:
          mask = mask & (dist_nn[..., 0] <= (ratio_thresh**2)*dist_nn[..., 1])
      if distance_thresh:
          mask = mask & (dist_nn[..., 0] <= distance_thresh**2)
      matches = torch.where(mask, ind_nn[..., 0], ind_nn.new_tensor(-1))
      scores = torch.where(mask, (sim_nn[..., 0]+1)/2, sim_nn.new_tensor(0))
      return matches, scores
  
  
  def mutual_check(m0, m1):
      inds0 = torch.arange(m0.shape[-1], device=m0.device)
      loop = torch.gather(m1, -1, torch.where(m0 > -1, m0, m0.new_tensor(0)))
      ok = (m0 > -1) & (inds0 == loop)
      m0_new = torch.where(ok, m0, m0.new_tensor(-1))
      return m0_new
  
  
  class NearestNeighbor(BaseModel):
      default_conf = {
          'ratio_threshold': None,
          'distance_threshold': None,
          'do_mutual_check': True,
      }
      required_inputs = ['descriptors0', 'descriptors1']
  
      def _init(self, conf):
          pass
  
      def _forward(self, data):
          if data['descriptors0'].size(-1) == 0 or data['descriptors1'].size(-1) == 0:
              matches0 = torch.full(
                  data['descriptors0'].shape[:2], -1,
                  device=data['descriptors0'].device)
              return {
                  'matches0': matches0,
                  'matching_scores0': torch.zeros_like(matches0)
              }
          ratio_threshold = self.conf['ratio_threshold']
          if data['descriptors0'].size(-1) == 1 or data['descriptors1'].size(-1) == 1:
              ratio_threshold = None
          sim = torch.einsum(
              'bdn,bdm->bnm', data['descriptors0'], data['descriptors1'])
          matches0, scores0 = find_nn(
              sim, ratio_threshold, self.conf['distance_threshold'])
          if self.conf['do_mutual_check']:
              matches1, scores1 = find_nn(
                  sim.transpose(1, 2), ratio_threshold,
                  self.conf['distance_threshold'])
              matches0 = mutual_check(matches0, matches1)
          return {
              'matches0': matches0,
              'matching_scores0': scores0,
          }
  