docstr: Function that computes 1 pass of non-offsetted Group-Wise correlation
docstr_tok: ["Function","that","computes","1","pass","of","non-offsetted","Group-Wise","correlation"]
code: def forward(
      self,
      left_feature: Tensor,
      right_feature: Tensor,
      flow: Tensor,
      window_type: str = "1d",
  ) -> Tensor:
      """Function that computes 1 pass of non-offsetted Group-Wise correlation"""
      coords = make_coords_grid(
          left_feature.shape[0],
          left_feature.shape[2],
          left_feature.shape[3],
          device=str(left_feature.device),
      )
  
      # we offset the coordinate grid in the flow direction
      coords = coords + flow
      coords = coords.permute(0, 2, 3, 1)
      # resample right features according to off-setted grid
      right_feature = grid_sample(
          right_feature, coords, mode="bilinear", align_corners=True
      )
  
      # use_small_patch is a flag by which we decide on how many axes
      # we perform candidate search. See section 3.1 ``Deformable search window`` & Figure 4 in the paper.
      patch_size_list = self.patch_sizes[window_type]
      dilate_size_list = self.dilate_sizes[window_type]
  
      # chunking the left and right feature to perform group-wise correlation
      # mechanism similar to GroupNorm. See section 3.1 ``Group-wise correlation``.
      left_groups = torch.chunk(left_feature, self.groups, dim=1)
      right_groups = torch.chunk(right_feature, self.groups, dim=1)
  
      correlations = []
      # this boils down to rather than performing the correlation product
      # over the entire C dimensions, we use subsets of C to get multiple correlation sets
      for i in range(len(patch_size_list)):
          correlation = get_correlation(
              left_groups[i],
              right_groups[i],
              patch_size_list[i],
              dilate_size_list[i],
          )
          correlations.append(correlation)
      final_correlations = torch.cat(correlations, dim=1)
      return final_correlations
  
imports: import math
  from functools import partial
  from typing import Callable, Dict, Iterable, List, Optional, Tuple
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import torchvision.models.optical_flow.raft
  from torch import Tensor
  from torchvision.models._api import register_model, Weights, WeightsEnum
  from torchvision.models._utils import handle_legacy_interface
  from torchvision.models.optical_flow._utils import grid_sample, make_coords_grid, upsample_flow
  from torchvision.ops import Conv2dNormActivation
  from torchvision.prototype.transforms._presets import StereoMatching
label_desc: Use image node_id 2 for calling the IterativeCorrelationLayer obj's underlying member method code with example usage: obj.forward(left_feature, right_feature, flow, window_type) and returns: final_correlations