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 forward(
      self,
      left_image: Tensor,
      right_image: Tensor,
      flow_init: Optional[Tensor] = None,
      num_iters: int = 10,
  ) -> List[Tensor]:
      features = torch.cat([left_image, right_image], dim=0)
      features = self.feature_encoder(features)
      left_features, right_features = features.chunk(2, dim=0)
  
      # update block network state and input context are derived from the left feature map
      net, ctx = left_features.chunk(2, dim=1)
      net = torch.tanh(net)
      ctx = torch.relu(ctx)
  
      # will output lists of tensor.
      l_pyramid = self.downsampling_pyramid(left_features)
      r_pyramid = self.downsampling_pyramid(right_features)
      net_pyramid = self.downsampling_pyramid(net)
      ctx_pyramid = self.downsampling_pyramid(ctx)
  
      # we store in reversed order because we process the pyramid from top to bottom
      l_pyramid = {
          res: l_pyramid[idx]
          for idx, res in enumerate(self.resolutions)
      }
      r_pyramid = {
          res: r_pyramid[idx]
          for idx, res in enumerate(self.resolutions)
      }
      net_pyramid = {
          res: net_pyramid[idx]
          for idx, res in enumerate(self.resolutions)
      }
      ctx_pyramid = {
          res: ctx_pyramid[idx]
          for idx, res in enumerate(self.resolutions)
      }
  
      # offsets for sampling pixel candidates in the correlation ops
      offsets: Dict[str, Tensor] = {}
      for resolution, offset_conv in self.offset_convs.items():
          feature_map = l_pyramid[resolution]
          offset = offset_conv(feature_map)
          offsets[resolution] = (torch.sigmoid(offset) - 0.5) * 2.0
  
      # the smallest resolution is prepared for passing through self attention
      min_res = self.resolutions[-1]
      max_res = self.resolutions[0]
  
      B, C, MIN_H, MIN_W = l_pyramid[min_res].shape
      # add positional encodings
      l_pyramid[min_res] = self.positional_encodings(l_pyramid[min_res])
      r_pyramid[min_res] = self.positional_encodings(r_pyramid[min_res])
      # reshaping for transformer
      l_pyramid[min_res] = (
          l_pyramid[min_res]
          .permute(0, 2, 3, 1)
          .reshape(B, MIN_H * MIN_W, C)
      )
      r_pyramid[min_res] = (
          r_pyramid[min_res]
          .permute(0, 2, 3, 1)
          .reshape(B, MIN_H * MIN_W, C)
      )
      # perform self attention
      l_pyramid[min_res], r_pyramid[min_res] = self.self_attn_block(
          l_pyramid[min_res], r_pyramid[min_res]
      )
      # now we need to reshape back into [B, C, H, W] format
      l_pyramid[min_res] = (
          l_pyramid[min_res]
          .reshape(B, MIN_H, MIN_W, C)
          .permute(0, 3, 1, 2)
      )
      r_pyramid[min_res] = (
          r_pyramid[min_res]
          .reshape(B, MIN_H, MIN_W, C)
          .permute(0, 3, 1, 2)
      )
  
      predictions: List[Tensor] = []
      flow_estimates: Dict[str, Tensor] = {}
      # we added this because of torch.script.jit
      # also, the predicition prior is always going to have the
      # spatial size of the features outputted by the feature encoder
      flow_pred_prior: Tensor = torch.empty(
          size=(B, 2, left_features.shape[2], left_features.shape[3]),
          dtype=l_pyramid[max_res].dtype,
          device=l_pyramid[max_res].device,
      )
  
      if flow_init is not None:
          scale = l_pyramid[max_res].shape[2] / flow_init.shape[2]
          # in CREStereo implementation they multiply with -scale instead of scale
          # this can be either a downsample or an upsample based on the cascaded inference
          # configuration
  
          # we use a -scale because the flow used inside the network is a negative flow
          # from the right to the left, so we flip the flow direction
          flow_estimates[max_res] = -scale * F.interpolate(
              input=flow_init,
              size=l_pyramid[max_res].shape[2:],
              mode="bilinear",
              align_corners=True,
          )
  
      # when not provided with a flow prior, we construct one using the lower resolution maps
      else:
          # initialize a zero flow with the smallest resolution
          flow = torch.zeros(
              size=(B, 2, MIN_H, MIN_W),
              device=left_features.device,
              dtype=left_features.dtype,
          )
  
          # flows from coarse resolutions are refined similarly
          # we always need to fetch the next pyramid feature map as well
          # when updating coarse resolutions, therefore we create a reversed
          # view which has its order synced with the ModuleDict keys iterator
          coarse_resolutions: List[str] = self.resolutions[
              ::-1
          ]  # using slicing because of torch.jit.script
          fine_grained_resolution = max_res
  
          # set the coarsest flow to the zero flow
          flow_estimates[coarse_resolutions[0]] = flow
  
          # the correlation layer ModuleDict will contain layers ordered from coarse to fine resolution
          # i.e ["1 / 16", "1 / 8", "1 / 4"]
          # the correlation layer ModuleDict has layers for all the resolutions except the fine one
          # i.e {"1 / 16": Module, "1 / 8": Module}
          # for these resolution we perform only half of the number of refinement iterations
          for idx, (resolution, correlation_layer) in enumerate(
              self.correlation_layers.items()
          ):
              # compute the scale difference between the first pyramid scale and the current pyramid scale
              scale_to_base = (
                  l_pyramid[fine_grained_resolution].shape[2]
                  // l_pyramid[resolution].shape[2]
              )
              for it in range(num_iters // 2):
                  # set whether we want to search on (X, Y) axes for correlation or just on X axis
                  window_type = self._get_window_type(it)
                  # we consider this a prior, therefore we do not want to back-propagate through it
                  flow_estimates[resolution] = flow_estimates[
                      resolution
                  ].detach()
  
                  correlations = correlation_layer(
                      l_pyramid[resolution],  # left
                      r_pyramid[resolution],  # right
                      flow_estimates[resolution],
                      offsets[resolution],
                      window_type,
                  )
  
                  # update the recurrent network state and the flow deltas
                  (
                      net_pyramid[resolution],
                      delta_flow,
                  ) = self.update_block(
                      net_pyramid[resolution],
                      ctx_pyramid[resolution],
                      correlations,
                      flow_estimates[resolution],
                  )
  
                  # the convex upsampling weights are computed w.r.t.
                  # the recurrent update state
                  up_mask = self.mask_predictor(net_pyramid[resolution])
                  flow_estimates[resolution] = (
                      flow_estimates[resolution] + delta_flow
                  )
                  # convex upsampling with the initial feature encoder downsampling rate
                  flow_pred_prior = upsample_flow(
                      flow_estimates[resolution],
                      up_mask,
                      factor=self.downsampling_factors[0],
                  )
                  # we then bilinear upsample to the final resolution
                  # we use a factor that's equivalent to the difference between
                  # the current downsample resolution and the base downsample resolution
                  #
                  # i.e. if a 1 / 16 flow is upsampled by 4 (base downsampling) we get a 1 / 4 flow.
                  # therefore we have to further upscale it by the difference between
                  # the current level 1 / 16 and the base level 1 / 4.
                  #
                  # we use a -scale because the flow used inside the network is a negative flow
                  # from the right to the left, so we flip the flow direction in order to get the
                  # left to right flow
                  flow_pred = -upsample_flow(
                      flow_pred_prior, None, factor=scale_to_base
                  )
                  predictions.append(flow_pred)
  
              # when constructing the next resolution prior, we resample w.r.t
              # to the scale of the next level in the pyramid
              next_resolution = coarse_resolutions[idx + 1]
              scale_to_next = (
                  l_pyramid[next_resolution].shape[2]
                  / flow_pred_prior.shape[2]
              )
              # we use the flow_up_prior because this is a more accurate estimation of the true flow
              # due to the convex upsample, which resembles a learned super-resolution module.
              # this is not necessarily an upsample, it can be a downsample, based on the provided configuration
              flow_estimates[
                  next_resolution
              ] = -scale_to_next * F.interpolate(
                  input=flow_pred_prior,
                  size=l_pyramid[next_resolution].shape[2:],
                  mode="bilinear",
                  align_corners=True,
              )
  
      # finally we will be doing a full pass through the fine-grained resolution
      # this coincides with the maximum resolution
  
      # we keep a separate loop here in order to avoid python control flow
      # to decide how many iterations should we do based on the current resolution
      # furthermore, if provided with an initial flow, there is no need to generate
      # a prior estimate when moving into the final refinement stage
  
      for it in range(num_iters):
          search_window_type = self._get_window_type(it)
  
          flow_estimates[max_res] = flow_estimates[max_res].detach()
          # we run the fine-grained resolution correlations in iterative mode
          # this means that we are using the fixed window pixel selections
          # instead of the deformed ones as with the previous steps
          correlations = self.max_res_correlation_layer(
              l_pyramid[max_res],
              r_pyramid[max_res],
              flow_estimates[max_res],
              extra_offset=None,
              window_type=search_window_type,
              iter_mode=True,
          )
  
          net_pyramid[max_res], delta_flow = self.update_block(
              net_pyramid[max_res],
              ctx_pyramid[max_res],
              correlations,
              flow_estimates[max_res],
          )
  
          up_mask = self.mask_predictor(net_pyramid[max_res])
          flow_estimates[max_res] = flow_estimates[max_res] + delta_flow
          # at the final resolution we simply do a convex upsample using the base downsample rate
          flow_pred = -upsample_flow(
              flow_estimates[max_res],
              up_mask,
              factor=self.downsampling_factors[0],
          )
          predictions.append(flow_pred)
  
      return predictions
  
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 3 for calling the CREStereo obj's underlying member method code with example usage: obj.forward(left_image, right_image, flow_init, num_iters) and returns: predictions