docstr: Function that computes 1 pass of offsetted Group-Wise correlation
  
  If the class was provided with an attention layer, the left and right feature maps
  will be passed through a transformer first
docstr_tok: ["Function","that","computes","1","pass","of","offsetted","Group-Wise","correlation","If","the","class","was","provided","with","an","attention","layer",",","the","left","and","right","feature","maps","will","be","passed","through","a","transformer","first"]
code: def forward(
      self,
      left_feature: Tensor,
      right_feature: Tensor,
      flow: Tensor,
      extra_offset: Tensor,
      window_type: str = "1d",
  ) -> Tensor:
      """Function that computes 1 pass of offsetted Group-Wise correlation
  
      If the class was provided with an attention layer, the left and right feature maps
      will be passed through a transformer first
      """
      B, C, H, W = left_feature.shape
  
      if self.attention_module is not None:
          # prepare for transformer required input shapes
          left_feature = left_feature.permute(0, 2, 3, 1).reshape(
              B, H * W, C
          )
          right_feature = right_feature.permute(0, 2, 3, 1).reshape(
              B, H * W, C
          )
          # this can be either self attention or cross attention, hence the tuple return
          left_feature, right_feature = self.attention_module(
              left_feature, right_feature
          )
          left_feature = left_feature.reshape(B, H, W, C).permute(
              0, 3, 1, 2
          )
          right_feature = right_feature.reshape(B, H, W, C).permute(
              0, 3, 1, 2
          )
  
      left_groups = torch.chunk(left_feature, self.groups, dim=1)
      right_groups = torch.chunk(right_feature, self.groups, dim=1)
  
      num_search_candidates = self.search_pixels
      # for each pixel (i, j) we have a number of search candidates
      # thus, for each candidate we should have an X-axis and Y-axis offset value
      extra_offset = extra_offset.reshape(
          B, num_search_candidates, 2, H, W
      ).permute(0, 1, 3, 4, 2)
  
      patch_size_list = self.patch_sizes[window_type]
      dilate_size_list = self.dilate_sizes[window_type]
  
      group_channels = C // self.groups
      correlations = []
  
      for i in range(len(patch_size_list)):
          left_group, right_group = left_groups[i], right_groups[i]
          patch_size, dilate = patch_size_list[i], dilate_size_list[i]
  
          di_y, di_x = dilate
          ps_y, ps_x = patch_size
          # define the search based on the window patch shape
          ry, rx = ps_y // 2 * di_y, ps_x // 2 * di_x
  
          # base offsets for search (i.e. where to look on the search index)
          x_grid, y_grid = torch.meshgrid(
              torch.arange(-rx, rx + 1, di_x),
              torch.arange(-ry, ry + 1, di_y),
              indexing="xy",
          )
          x_grid, y_grid = x_grid.to(flow.device), y_grid.to(
              flow.device
          )
          offsets = torch.stack((x_grid, y_grid))
          offsets = offsets.reshape(2, -1).permute(1, 0)
  
          for d in (0, 2, 3):
              offsets = offsets.unsqueeze(d)
          # extra offsets for search (i.e. deformed search indexes. Similar concept to deformable convolutions)
          offsets = offsets + extra_offset
  
          coords = (
              make_coords_grid(
                  left_feature.shape[0],
                  left_feature.shape[2],
                  left_feature.shape[3],
                  device=str(left_feature.device),
              )
              + flow
          )
          coords = coords.permute(0, 2, 3, 1).unsqueeze(1)
          coords = coords + offsets
          coords = coords.reshape(B, -1, W, 2)
  
          right_group = grid_sample(
              right_group, coords, mode="bilinear", align_corners=True
          )
          # we do not need to perform any window shifting because the grid sample op
          # will return a multi-view right based on the num_search_candidates dimension in the offsets
          right_group = right_group.reshape(B, group_channels, -1, H, W)
          left_group = left_group.reshape(B, group_channels, -1, H, W)
          correlation = torch.mean(left_group * right_group, dim=1)
          correlations.append(correlation)
  
      final_correlation = torch.cat(correlations, dim=1)
      return final_correlation
  
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 AttentionOffsetCorrelationLayer obj's underlying member method code with example usage: obj.forward(left_feature, right_feature, flow, extra_offset, window_type) and returns: final_correlation