code: def _reference_vicregl_vicreg_loss(
      x: Tensor,
      y: Tensor,
      inv_coeff: float = 25.0,
      var_coeff: float = 25.0,
      cov_coeff: float = 1.0,
  ) -> Tensor:
      # Loss implementation from VICRegL:
      # https://github.com/facebookresearch/VICRegL/blob/803ae4c8cd1649a820f03afb4793763e95317620/main_vicregl.py#L284
      repr_loss = inv_coeff * F.mse_loss(x, y)
  
      x = x - x.mean(0)
      y = y - y.mean(0)
  
      std_x = torch.sqrt(x.var(dim=0) + 0.0001)
      std_y = torch.sqrt(y.var(dim=0) + 0.0001)
      std_loss = var_coeff * (
          torch.mean(F.relu(1.0 - std_x)) / 2
          + torch.mean(F.relu(1.0 - std_y)) / 2
      )
  
      x = x.permute((1, 0, 2))
      y = y.permute((1, 0, 2))
  
      *_, sample_size, num_channels = x.shape
      non_diag_mask = ~torch.eye(
          num_channels, device=x.device, dtype=torch.bool
      )
      # Center features
      # centered.shape = NC
      x = x - x.mean(dim=-2, keepdim=True)
      y = y - y.mean(dim=-2, keepdim=True)
  
      cov_x = torch.einsum("...nc,...nd->...cd", x, x) / (
          sample_size - 1
      )
      cov_y = torch.einsum("...nc,...nd->...cd", y, y) / (
          sample_size - 1
      )
      cov_loss = (
          cov_x[..., non_diag_mask].pow(2).sum(-1) / num_channels
      ) / 2 + (
          cov_y[..., non_diag_mask].pow(2).sum(-1) / num_channels
      ) / 2
      cov_loss = cov_loss.mean()
      cov_loss = cov_coeff * cov_loss
  
      return repr_loss + std_loss + cov_loss
  
imports: import unittest
  import pytest
  import torch
  import torch.nn.functional
  from pytest_mock import MockerFixture
  from torch import Tensor
  from torch import distributed
  from lightly.loss import VICRegLoss
label_desc: Use image node_id 2 for calling a global function with example usage: _reference_vicregl_vicreg_loss(x, y, inv_coeff, var_coeff, cov_coeff) and returns: unknown