class_docstr: The correlation block.
  
  Creates a correlation pyramid with ``num_levels`` levels from the outputs of the feature encoder,
  and then indexes from this pyramid to create correlation features.
  The "indexing" of a given centroid pixel x' is done by concatenating its surrounding neighbors that
  are within a ``radius``, according to the infinity norm (see paper section 3.2).
  Note: typo in the paper, it should be infinity norm, not 1-norm.
class_docstr_tok: ["The","correlation","block",".","Creates","a","correlation","pyramid","with","``","num_levels","``","levels","from","the","outputs","of","the","feature","encoder",",","and","then","indexes","from","this","pyramid","to","create","correlation","features",".","The","``","indexing","''","of","a","given","centroid","pixel","x","'","is","done","by","concatenating","its","surrounding","neighbors","that","are","within","a","``","radius","``",",","according","to","the","infinity","norm","(","see","paper","section","3.2",")",".","Note",":","typo","in","the","paper",",","it","should","be","infinity","norm",",","not","1-norm","."]
docstr: Return correlation features by indexing from the pyramid.
docstr_tok: ["Return","correlation","features","by","indexing","from","the","pyramid","."]
code: def index_pyramid(self, centroids_coords):
      """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
      )
      dj = torch.linspace(
          -self.radius, self.radius, neighborhood_side_len
      )
      delta = torch.stack(
          torch.meshgrid(di, dj, indexing="ij"), dim=-1
      ).to(centroids_coords.device)
      delta = delta.view(
          1, neighborhood_side_len, neighborhood_side_len, 2
      )
  
      batch_size, _, h, w = centroids_coords.shape  # _ = 2
      centroids_coords = centroids_coords.permute(0, 2, 3, 1).reshape(
          batch_size * h * w, 1, 1, 2
      )
  
      indexed_pyramid = []
      for corr_volume in self.corr_pyramid:
          sampling_coords = (
              centroids_coords + delta
          )  # end shape is (batch_size * h * w, side_len, side_len, 2)
          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)
      if corr_features.shape != expected_output_shape:
          raise ValueError(
              f"Output shape of index pyramid is incorrect. Should be {expected_output_shape}, got {corr_features.shape}"
          )
  
      return corr_features
  
imports: from typing import List, Optional
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torch.nn.modules.batchnorm import BatchNorm2d
  from torch.nn.modules.instancenorm import InstanceNorm2d
  from torchvision.ops import Conv2dNormActivation
  from ...transforms._presets import OpticalFlow
  from ...utils import _log_api_usage_once
  from .._api import register_model, Weights, WeightsEnum
  from .._utils import handle_legacy_interface
  from ._utils import grid_sample, make_coords_grid, upsample_flow
label_desc: Use image node_id 3 for calling the CorrBlock obj's underlying member method code with example usage: obj.index_pyramid(centroids_coords) and returns: corr_features