class_docstr: Implements CREStereo from the `"Practical Stereo Matching via Cascaded Recurrent Network
  With Adaptive Correlation" <https://openaccess.thecvf.com/content/CVPR2022/papers/Li_Practical_Stereo_Matching_via_Cascaded_Recurrent_Network_With_Adaptive_Correlation_CVPR_2022_paper.pdf>`_ paper.
  Args:
      feature_encoder (raft.FeatureEncoder): Raft-like Feature Encoder module extract low-level features from inputs.
      update_block (raft.UpdateBlock): Raft-like Update Block which recursively refines a flow-map.
      flow_head (raft.FlowHead): Raft-like Flow Head which predics a flow-map from some inputs.
      self_attn_block (LocalFeatureTransformer): A Local Feature Transformer that performs self attention on the two feature maps.
      cross_attn_block (LocalFeatureTransformer): A Local Feature Transformer that performs cross attention between the two feature maps
          used in the Adaptive Group Correlation module.
      feature_downsample_rates (List[int]): The downsample rates used to build a feature pyramid from the outputs of the `feature_encoder`. Default: [2, 4]
      correlation_groups (int): In how many groups should the features be split when computer per-pixel correlation. Defaults 4.
      search_window_1d (Tuple[int, int]): The alternate search window size in the x and y directions for the 1D case. Defaults to (1, 9).
      search_dilate_1d (Tuple[int, int]): The dilation used in the `search_window_1d` when selecting pixels. Similar to `nn.Conv2d` dilate. Defaults to (1, 1).
      search_window_2d (Tuple[int, int]): The alternate search window size in the x and y directions for the 2D case. Defaults to (3, 3).
      search_dilate_2d (Tuple[int, int]): The dilation used in the `search_window_2d` when selecting pixels. Similar to `nn.Conv2d` dilate. Defaults to (1, 1).
class_docstr_tok: ["Implements","CREStereo","from","the","`","``","Practical","Stereo","Matching","via","Cascaded","Recurrent","Network","With","Adaptive","Correlation","''","<","https",":","\/\/openaccess.thecvf.com\/content\/CVPR2022\/papers\/Li_Practical_Stereo_Matching_via_Cascaded_Recurrent_Network_With_Adaptive_Correlation_CVPR_2022_paper.pdf",">","`","_","paper",".","Args",":","feature_encoder","(","raft.FeatureEncoder",")",":","Raft-like","Feature","Encoder","module","extract","low-level","features","from","inputs",".","update_block","(","raft.UpdateBlock",")",":","Raft-like","Update","Block","which","recursively","refines","a","flow-map",".","flow_head","(","raft.FlowHead",")",":","Raft-like","Flow","Head","which","predics","a","flow-map","from","some","inputs",".","self_attn_block","(","LocalFeatureTransformer",")",":","A","Local","Feature","Transformer","that","performs","self","attention","on","the","two","feature","maps",".","cross_attn_block","(","LocalFeatureTransformer",")",":","A","Local","Feature","Transformer","that","performs","cross","attention","between","the","two","feature","maps","used","in","the","Adaptive","Group","Correlation","module",".","feature_downsample_rates","(","List","[","int","]",")",":","The","downsample","rates","used","to","build","a","feature","pyramid","from","the","outputs","of","the","`","feature_encoder","`",".","Default",":","[","2",",","4","]","correlation_groups","(","int",")",":","In","how","many","groups","should","the","features","be","split","when","computer","per-pixel","correlation",".","Defaults","4.","search_window_1d","(","Tuple","[","int",",","int","]",")",":","The","alternate","search","window","size","in","the","x","and","y","directions","for","the","1D","case",".","Defaults","to","(","1",",","9",")",".","search_dilate_1d","(","Tuple","[","int",",","int","]",")",":","The","dilation","used","in","the","`","search_window_1d","`","when","selecting","pixels",".","Similar","to","`","nn.Conv2d","`","dilate",".","Defaults","to","(","1",",","1",")",".","search_window_2d","(","Tuple","[","int",",","int","]",")",":","The","alternate","search","window","size","in","the","x","and","y","directions","for","the","2D","case",".","Defaults","to","(","3",",","3",")",".","search_dilate_2d","(","Tuple","[","int",",","int","]",")",":","The","dilation","used","in","the","`","search_window_2d","`","when","selecting","pixels",".","Similar","to","`","nn.Conv2d","`","dilate",".","Defaults","to","(","1",",","1",")","."]
code: def __init__(
      self,
      *,
      feature_encoder: raft.FeatureEncoder,
      update_block: raft.UpdateBlock,
      flow_head: raft.FlowHead,
      self_attn_block: LocalFeatureTransformer,
      cross_attn_block: LocalFeatureTransformer,
      feature_downsample_rates: Tuple[int, ...] = (2, 4),
      correlation_groups: int = 4,
      search_window_1d: Tuple[int, int] = (1, 9),
      search_dilate_1d: Tuple[int, int] = (1, 1),
      search_window_2d: Tuple[int, int] = (3, 3),
      search_dilate_2d: Tuple[int, int] = (1, 1),
  ) -> None:
      super().__init__()
      self.output_channels = 2
  
      self.feature_encoder = feature_encoder
      self.update_block = update_block
      self.flow_head = flow_head
      self.self_attn_block = self_attn_block
  
      # average pooling for the feature encoder outputs
      self.downsampling_pyramid = PyramidDownsample(
          feature_downsample_rates
      )
      self.downsampling_factors: List[int] = [
          feature_encoder.downsample_factor
      ]
      base_downsample_factor: int = self.downsampling_factors[0]
      for rate in feature_downsample_rates:
          self.downsampling_factors.append(
              base_downsample_factor * rate
          )
  
      # output resolution tracking
      self.resolutions: List[str] = [
          f"1 / {factor}" for factor in self.downsampling_factors
      ]
      self.search_pixels = int(np.prod(search_window_1d))
  
      # flow convex upsampling mask predictor
      self.mask_predictor = ConvexMaskPredictor(
          in_channels=feature_encoder.output_dim // 2,
          hidden_size=feature_encoder.output_dim,
          upsample_factor=feature_encoder.downsample_factor,
          multiplier=0.25,
      )
  
      # offsets modules for offsetted feature selection
      self.offset_convs = nn.ModuleDict()
      self.correlation_layers = nn.ModuleDict()
  
      offset_conv_layer = partial(
          Conv2dNormActivation,
          in_channels=feature_encoder.output_dim,
          out_channels=self.search_pixels * 2,
          norm_layer=None,
          activation_layer=None,
      )
  
      # populate the dicts in top to bottom order
      # useful for iterating through torch.jit.script module given the network forward pass
      #
      # Ignore the largest resolution. We handle that separately due to torch.jit.script
      # not being able to access to runtime generated keys in ModuleDicts.
      # This way, we can keep a generic way of processing all pyramid levels but except
      # the final one
      iterative_correlation_layer = partial(
          IterativeCorrelationLayer,
          groups=correlation_groups,
          search_window_1d=search_window_1d,
          search_dilate_1d=search_dilate_1d,
          search_window_2d=search_window_2d,
          search_dilate_2d=search_dilate_2d,
      )
  
      attention_offset_correlation_layer = partial(
          AttentionOffsetCorrelationLayer,
          groups=correlation_groups,
          search_window_1d=search_window_1d,
          search_dilate_1d=search_dilate_1d,
          search_window_2d=search_window_2d,
          search_dilate_2d=search_dilate_2d,
      )
  
      for idx, resolution in enumerate(reversed(self.resolutions[1:])):
          # the largest resolution does use offset convolutions for sampling grid coords
          offset_conv = (
              None
              if idx == len(self.resolutions) - 1
              else offset_conv_layer()
          )
          if offset_conv:
              self.offset_convs[resolution] = offset_conv
              # only the lowest resolution uses the cross attention module when computing correlation scores
              attention_module = cross_attn_block if idx == 0 else None
              self.correlation_layers[
                  resolution
              ] = AdaptiveGroupCorrelationLayer(
                  iterative_correlation_layer=iterative_correlation_layer(),
                  attention_offset_correlation_layer=attention_offset_correlation_layer(
                      attention_module=attention_module
                  ),
              )
  
      # correlation layer for the largest resolution
      self.max_res_correlation_layer = AdaptiveGroupCorrelationLayer(
          iterative_correlation_layer=iterative_correlation_layer(),
          attention_offset_correlation_layer=attention_offset_correlation_layer(),
      )
  
      # simple 2D Postional Encodings
      self.positional_encodings = PositionalEncodingSine(
          feature_encoder.output_dim
      )
  
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 1 to create a new CREStereo object from inherited base classes: nn with example: obj = CREStereo()