class_docstr: The row-wise correlation block.
  
  Use indexes from correlation pyramid to create correlation features.
  The "indexing" of a given centroid pixel x' is done by concatenating its surrounding row neighbours
  within radius
class_docstr_tok: ["The","row-wise","correlation","block",".","Use","indexes","from","correlation","pyramid","to","create","correlation","features",".","The","``","indexing","''","of","a","given","centroid","pixel","x","'","is","done","by","concatenating","its","surrounding","row","neighbours","within","radius"]
docstr: Return correlation features by indexing from the pyramid.
code: def forward(
      self, centroids_coords: Tensor, corr_pyramid: List[Tensor]
  ) -> Tensor:
      """Return correlation features by indexing from the pyramid."""
      neighborhood_side_len = (
          2 * self.radius + 1
      )  # see note in __init__ about out_channels
      di = torch.linspace(
          -self.radius,
          self.radius,
          neighborhood_side_len,
          device=centroids_coords.device,
      )
      di = di.view(1, 1, neighborhood_side_len, 1).to(
          centroids_coords.device
      )
  
      (
          batch_size,
          _,
          h,
          w,
      ) = centroids_coords.shape  # _ = 2 but we only use the first one
      # We only consider 1d and take the first dim only
      centroids_coords = (
          centroids_coords[:, :1]
          .permute(0, 2, 3, 1)
          .reshape(batch_size * h * w, 1, 1, 1)
      )
  
      indexed_pyramid = []
      for corr_volume in corr_pyramid:
          x0 = (
              centroids_coords + di
          )  # end shape is (batch_size * h * w, 1, side_len, 1)
          y0 = torch.zeros_like(x0)
          sampling_coords = torch.cat([x0, y0], dim=-1)
          indexed_corr_volume = grid_sample(
              corr_volume,
              sampling_coords,
              align_corners=True,
              mode="bilinear",
          ).view(batch_size, h, w, -1)
          indexed_pyramid.append(indexed_corr_volume)
          centroids_coords = centroids_coords / 2
  
      corr_features = (
          torch.cat(indexed_pyramid, dim=-1)
          .permute(0, 3, 1, 2)
          .contiguous()
      )
  
      expected_output_shape = (batch_size, self.out_channels, h, w)
      torch._assert(
          corr_features.shape == expected_output_shape,
          f"Output shape of index pyramid is incorrect. Should be {expected_output_shape}, got {corr_features.shape}",
      )
      return corr_features
  
imports: from functools import partial
  from typing import Callable, List, Optional, Tuple
  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.models.optical_flow.raft import FlowHead, MotionEncoder, ResidualBlock
  from torchvision.ops import Conv2dNormActivation
  from torchvision.prototype.transforms._presets import StereoMatching
  from torchvision.utils import _log_api_usage_once
label_desc: Use image node_id 2 for calling the CorrBlock1d obj's underlying member method code with example usage: obj.forward(centroids_coords, corr_pyramid) and returns: corr_features