class_docstr: Implementation of the VICRegL loss from VICRegL paper [0].
  
  This implementation follows the code published by the authors [1].
  
  - [0]: VICRegL, 2022, https://arxiv.org/abs/2210.01571
  - [1]: https://github.com/facebookresearch/VICRegL
  
  Attributes:
      lambda_param:
          Coefficient for the invariance term of the loss.
      mu_param:
          Coefficient for the variance term of the loss.
      nu_param:
          Coefficient for the covariance term of the loss.
      alpha:
          Coefficient to weight global with local loss. The final loss is computed as
          (self.alpha * global_loss + (1-self.alpha) * local_loss).
      gather_distributed:
          If True then the cross-correlation matrices from all gpus are gathered and
          summed before the loss calculation.
      eps:
          Epsilon for numerical stability.
      num_matches:
          Number of local features to match using nearest neighbors.
  
  Examples:
  
      >>> # initialize loss function
      >>> criterion = VICRegLLoss()
      >>> transform = VICRegLTransform(n_global_views=2, n_local_views=4)
      >>>
      >>> # generate two random transforms of images
      >>> views_and_grids = transform(images)
      >>> views = views_and_grids[:6] # 2 global views + 4 local views
      >>> grids = views_and_grids[6:]
      >>>
      >>> # feed through model images
      >>> features = [model(view) for view in views]
      >>>
      >>> # calculate loss
      >>> loss = criterion(
      ...     global_view_features=features[:2],
      ...     global_view_grids=grids[:2],
      ...     local_view_features=features[2:],
      ...     local_view_grids=grids[2:],
      ... )
class_docstr_tok: ["Implementation","of","the","VICRegL","loss","from","VICRegL","paper","[","0","]",".","This","implementation","follows","the","code","published","by","the","authors","[","1","]",".","-","[","0","]",":","VICRegL",",","2022",",","https",":","\/\/arxiv.org\/abs\/2210.01571","-","[","1","]",":","https",":","\/\/github.com\/facebookresearch\/VICRegL","Attributes",":","lambda_param",":","Coefficient","for","the","invariance","term","of","the","loss",".","mu_param",":","Coefficient","for","the","variance","term","of","the","loss",".","nu_param",":","Coefficient","for","the","covariance","term","of","the","loss",".","alpha",":","Coefficient","to","weight","global","with","local","loss",".","The","final","loss","is","computed","as","(","self.alpha","*","global_loss","+","(","1-self.alpha",")","*","local_loss",")",".","gather_distributed",":","If","True","then","the","cross-correlation","matrices","from","all","gpus","are","gathered","and","summed","before","the","loss","calculation",".","eps",":","Epsilon","for","numerical","stability",".","num_matches",":","Number","of","local","features","to","match","using","nearest","neighbors",".","Examples",":",">",">",">","#","initialize","loss","function",">",">",">","criterion","=","VICRegLLoss","(",")",">",">",">","transform","=","VICRegLTransform","(","n_global_views=2",",","n_local_views=4",")",">",">",">",">",">",">","#","generate","two","random","transforms","of","images",">",">",">","views_and_grids","=","transform","(","images",")",">",">",">","views","=","views_and_grids","[",":6","]","#","2","global","views","+","4","local","views",">",">",">","grids","=","views_and_grids","[","6",":","]",">",">",">",">",">",">","#","feed","through","model","images",">",">",">","features","=","[","model","(","view",")","for","view","in","views","]",">",">",">",">",">",">","#","calculate","loss",">",">",">","loss","=","criterion","(","...","global_view_features=features","[",":2","]",",","...","global_view_grids=grids","[",":2","]",",","...","local_view_features=features","[","2",":","]",",","...","local_view_grids=grids","[","2",":","]",",","...",")"]
docstr: Finds num_matches closest neighbors of input_features in candidate_features
  based on the distance between the features defined by input_grid and
  candidate_grid.
  
  Args:
      input_features:
          Local features tensor with shape (batch_size, height * width, dim).
      candidate_features:
          Local features tensor with shape (batch_size, height * width, dim).
          Note that height and width can be different from input_features.
      input_grid:
          Grid tensor with shape (batch_size, height, width, 2).
      candidate_grid:
          Grid tensor with shape (batch_size, height, width, 2). Note that height
          and width can be different from input_grid.
  
  Returns:
      (nn_input, nn_candidate) tuple containing two tensors with shape
      (batch_size, num_matches, dim).
docstr_tok: ["Finds","num_matches","closest","neighbors","of","input_features","in","candidate_features","based","on","the","distance","between","the","features","defined","by","input_grid","and","candidate_grid",".","Args",":","input_features",":","Local","features","tensor","with","shape","(","batch_size",",","height","*","width",",","dim",")",".","candidate_features",":","Local","features","tensor","with","shape","(","batch_size",",","height","*","width",",","dim",")",".","Note","that","height","and","width","can","be","different","from","input_features",".","input_grid",":","Grid","tensor","with","shape","(","batch_size",",","height",",","width",",","2",")",".","candidate_grid",":","Grid","tensor","with","shape","(","batch_size",",","height",",","width",",","2",")",".","Note","that","height","and","width","can","be","different","from","input_grid",".","Returns",":","(","nn_input",",","nn_candidate",")","tuple","containing","two","tensors","with","shape","(","batch_size",",","num_matches",",","dim",")","."]
code: def _nearest_neighbors_on_grid(
      self,
      input_features: Tensor,
      candidate_features: Tensor,
      input_grid: Tensor,
      candidate_grid: Tensor,
      num_matches: int,
  ) -> Tuple[Tensor, Tensor]:
      """Finds num_matches closest neighbors of input_features in candidate_features
      based on the distance between the features defined by input_grid and
      candidate_grid.
  
      Args:
          input_features:
              Local features tensor with shape (batch_size, height * width, dim).
          candidate_features:
              Local features tensor with shape (batch_size, height * width, dim).
              Note that height and width can be different from input_features.
          input_grid:
              Grid tensor with shape (batch_size, height, width, 2).
          candidate_grid:
              Grid tensor with shape (batch_size, height, width, 2). Note that height
              and width can be different from input_grid.
  
      Returns:
          (nn_input, nn_candidate) tuple containing two tensors with shape
          (batch_size, num_matches, dim).
      """
      distances = torch.cdist(input_grid, candidate_grid)
      return nearest_neighbors(
          input_features, candidate_features, distances, num_matches
      )
  
file_path: lightly/lightly/loss/vicregl_loss.py
imports: from typing import Optional, Sequence, Tuple
  import torch
  import torch.distributed
  from torch import Tensor
  from lightly.loss.vicreg_loss import VICRegLoss, covariance_loss, invariance_loss, variance_loss
  from lightly.models.utils import nearest_neighbors
  from lightly.utils.dist import gather
label_desc: Use image node_id 10 for calling the VICRegLLoss obj's underlying member method code with example usage: obj._nearest_neighbors_on_grid(input_features, candidate_features, input_grid, candidate_grid, num_matches) and returns: nearest_neighbors