code: def _crestereo(
      *,
      weights: Optional[WeightsEnum],
      progress: bool,
      # Feature Encoder
      feature_encoder_layers: Tuple[int, int, int, int, int],
      feature_encoder_strides: Tuple[int, int, int, int],
      feature_encoder_block: Callable[..., nn.Module],
      feature_encoder_norm_layer: Callable[..., nn.Module],
      # Average Pooling Pyramid
      feature_downsample_rates: Tuple[int, ...],
      # Adaptive Correlation Layer
      corr_groups: int,
      corr_search_window_2d: Tuple[int, int],
      corr_search_dilate_2d: Tuple[int, int],
      corr_search_window_1d: Tuple[int, int],
      corr_search_dilate_1d: Tuple[int, int],
      # Flow head
      flow_head_hidden_size: int,
      # Recurrent block
      recurrent_block_hidden_state_size: int,
      recurrent_block_kernel_size: Tuple[
          Tuple[int, int], Tuple[int, int]
      ],
      recurrent_block_padding: Tuple[Tuple[int, int], Tuple[int, int]],
      # Motion Encoder
      motion_encoder_corr_layers: Tuple[int, int],
      motion_encoder_flow_layers: Tuple[int, int],
      motion_encoder_out_channels: int,
      # Transformer Blocks
      num_attention_heads: int,
      num_self_attention_layers: int,
      num_cross_attention_layers: int,
      self_attention_module: Callable[..., nn.Module],
      cross_attention_module: Callable[..., nn.Module],
      **kwargs,
  ) -> CREStereo:
      feature_encoder = kwargs.pop(
          "feature_encoder", None
      ) or raft.FeatureEncoder(
          block=feature_encoder_block,
          layers=feature_encoder_layers,
          strides=feature_encoder_strides,
          norm_layer=feature_encoder_norm_layer,
      )
  
      if feature_encoder.output_dim % corr_groups != 0:
          raise ValueError(
              f"Final ``feature_encoder_layers`` size should be divisible by ``corr_groups`` argument."
              f"Feature encoder output size : {feature_encoder.output_dim}, Correlation groups: {corr_groups}."
          )
  
      motion_encoder = kwargs.pop(
          "motion_encoder", None
      ) or raft.MotionEncoder(
          in_channels_corr=corr_groups
          * int(np.prod(corr_search_window_1d)),
          corr_layers=motion_encoder_corr_layers,
          flow_layers=motion_encoder_flow_layers,
          out_channels=motion_encoder_out_channels,
      )
  
      out_channels_context = (
          feature_encoder_layers[-1] - recurrent_block_hidden_state_size
      )
      recurrent_block = kwargs.pop(
          "recurrent_block", None
      ) or raft.RecurrentBlock(
          input_size=motion_encoder.out_channels + out_channels_context,
          hidden_size=recurrent_block_hidden_state_size,
          kernel_size=recurrent_block_kernel_size,
          padding=recurrent_block_padding,
      )
  
      flow_head = kwargs.pop("flow_head", None) or raft.FlowHead(
          in_channels=out_channels_context,
          hidden_size=flow_head_hidden_size,
      )
  
      update_block = raft.UpdateBlock(
          motion_encoder=motion_encoder,
          recurrent_block=recurrent_block,
          flow_head=flow_head,
      )
  
      self_attention_module = (
          kwargs.pop("self_attention_module", None) or LinearAttention
      )
      self_attn_block = LocalFeatureTransformer(
          dim_model=feature_encoder.output_dim,
          num_heads=num_attention_heads,
          attention_directions=["self"] * num_self_attention_layers,
          attention_module=self_attention_module,
      )
  
      cross_attention_module = (
          kwargs.pop("cross_attention_module", None) or LinearAttention
      )
      cross_attn_block = LocalFeatureTransformer(
          dim_model=feature_encoder.output_dim,
          num_heads=num_attention_heads,
          attention_directions=["cross"] * num_cross_attention_layers,
          attention_module=cross_attention_module,
      )
  
      model = CREStereo(
          feature_encoder=feature_encoder,
          update_block=update_block,
          flow_head=flow_head,
          self_attn_block=self_attn_block,
          cross_attn_block=cross_attn_block,
          feature_downsample_rates=feature_downsample_rates,
          correlation_groups=corr_groups,
          search_window_1d=corr_search_window_1d,
          search_window_2d=corr_search_window_2d,
          search_dilate_1d=corr_search_dilate_1d,
          search_dilate_2d=corr_search_dilate_2d,
      )
  
      if weights is not None:
          model.load_state_dict(
              weights.get_state_dict(progress=progress, check_hash=True)
          )
  
      return model
  
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 4 for calling a global function with example usage: _crestereo() and returns: model