docstr: RAFT-Stereo model from
  `RAFT-Stereo: Multilevel Recurrent Field Transforms for Stereo Matching <https://arxiv.org/abs/2109.07547>`_.
  
  args:
      feature_encoder (FeatureEncoder): The feature encoder. Its input is the concatenation of ``left_image`` and ``right_image``.
      context_encoder (MultiLevelContextEncoder): The context encoder. Its input is ``left_image``.
          It has multi-level output and each level will have 2 parts:
  
          - one part will be used as the actual "context", passed to the recurrent unit of the ``update_block``
          - one part will be used to initialize the hidden state of the recurrent unit of
            the ``update_block``
  
      corr_pyramid (CorrPyramid1d): Module to build the correlation pyramid from feature encoder output
      corr_block (CorrBlock1d): The correlation block, which uses the correlation pyramid indexes
          to create correlation features. It takes the coordinate of the centroid pixel and correlation pyramid
          as input and returns the correlation features.
          It must expose an ``out_channels`` attribute.
  
      update_block (MultiLevelUpdateBlock): The update block, which contains the motion encoder, and the recurrent unit.
          It takes as input the hidden state of its recurrent unit, the context, the correlation
          features, and the current predicted disparity. It outputs an updated hidden state
      disparity_head (nn.Module): The disparity head block will convert from the hidden state into changes in disparity.
      mask_predictor (nn.Module, optional): Predicts the mask that will be used to upsample the predicted flow.
          If ``None`` (default), the flow is upsampled using interpolation.
      slow_fast (bool): A boolean that specify whether we should use slow-fast GRU or not. See RAFT-Stereo paper
          on section 3.4 for more detail.
code: def __init__(
      self,
      *,
      feature_encoder: FeatureEncoder,
      context_encoder: MultiLevelContextEncoder,
      corr_pyramid: CorrPyramid1d,
      corr_block: CorrBlock1d,
      update_block: MultiLevelUpdateBlock,
      disparity_head: nn.Module,
      mask_predictor: Optional[nn.Module] = None,
      slow_fast: bool = False,
  ):
      """RAFT-Stereo model from
      `RAFT-Stereo: Multilevel Recurrent Field Transforms for Stereo Matching <https://arxiv.org/abs/2109.07547>`_.
  
      args:
          feature_encoder (FeatureEncoder): The feature encoder. Its input is the concatenation of ``left_image`` and ``right_image``.
          context_encoder (MultiLevelContextEncoder): The context encoder. Its input is ``left_image``.
              It has multi-level output and each level will have 2 parts:
  
              - one part will be used as the actual "context", passed to the recurrent unit of the ``update_block``
              - one part will be used to initialize the hidden state of the recurrent unit of
                the ``update_block``
  
          corr_pyramid (CorrPyramid1d): Module to build the correlation pyramid from feature encoder output
          corr_block (CorrBlock1d): The correlation block, which uses the correlation pyramid indexes
              to create correlation features. It takes the coordinate of the centroid pixel and correlation pyramid
              as input and returns the correlation features.
              It must expose an ``out_channels`` attribute.
  
          update_block (MultiLevelUpdateBlock): The update block, which contains the motion encoder, and the recurrent unit.
              It takes as input the hidden state of its recurrent unit, the context, the correlation
              features, and the current predicted disparity. It outputs an updated hidden state
          disparity_head (nn.Module): The disparity head block will convert from the hidden state into changes in disparity.
          mask_predictor (nn.Module, optional): Predicts the mask that will be used to upsample the predicted flow.
              If ``None`` (default), the flow is upsampled using interpolation.
          slow_fast (bool): A boolean that specify whether we should use slow-fast GRU or not. See RAFT-Stereo paper
              on section 3.4 for more detail.
      """
      super().__init__()
      _log_api_usage_once(self)
  
      # This indicates that the disparity output will be only have 1 channel (represent horizontal axis).
      # We need this because some stereo matching model like CREStereo might have 2 channel on the output
      self.output_channels = 1
  
      self.feature_encoder = feature_encoder
      self.context_encoder = context_encoder
  
      self.base_downsampling_ratio = (
          feature_encoder.base_downsampling_ratio
      )
      self.num_level = self.context_encoder.num_level
      self.corr_pyramid = corr_pyramid
      self.corr_block = corr_block
      self.update_block = update_block
      self.disparity_head = disparity_head
      self.mask_predictor = mask_predictor
  
      hidden_dims = self.update_block.hidden_dims
      # Follow the original implementation to do pre convolution on the context
      # See: https://github.com/princeton-vl/RAFT-Stereo/blob/main/core/raft_stereo.py#L32
      self.context_convs = nn.ModuleList(
          [
              nn.Conv2d(
                  hidden_dims[i],
                  hidden_dims[i] * 3,
                  kernel_size=3,
                  padding=1,
              )
              for i in range(self.num_level)
          ]
      )
      self.slow_fast = slow_fast
  
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 1 to create a new RaftStereo object from inherited base classes: nn with example: obj = RaftStereo()