class_docstr: Reversible model builder. This builds the reversible transformer encoder
  and allows reversible training.
  
  Karttikeya Mangalam, Haoqi Fan, Yanghao Li, Chao-Yuan Wu, Bo Xiong,
  Christoph Feichtenhofer, Jitendra Malik
  "Reversible Vision Transformers"
  
  https://openaccess.thecvf.com/content/CVPR2022/papers/Mangalam_Reversible_Vision_Transformers_CVPR_2022_paper.pdf
class_docstr_tok: ["Reversible","model","builder",".","This","builds","the","reversible","transformer","encoder","and","allows","reversible","training",".","Karttikeya","Mangalam",",","Haoqi","Fan",",","Yanghao","Li",",","Chao-Yuan","Wu",",","Bo","Xiong",",","Christoph","Feichtenhofer",",","Jitendra","Malik","''","Reversible","Vision","Transformers","''","https",":","\/\/openaccess.thecvf.com\/content\/CVPR2022\/papers\/Mangalam_Reversible_Vision_Transformers_CVPR_2022_paper.pdf"]
docstr: The `__init__` method of any subclass should also contain these
      arguments.
  Args:
      cfg (CfgNode): model building configs, details are in the
          comments of the config file.
      model (nn.Module): parent MViT module this module forms
          a reversible encoder in.
docstr_tok: ["The","`","__init__","`","method","of","any","subclass","should","also","contain","these","arguments",".","Args",":","cfg","(","CfgNode",")",":","model","building","configs",",","details","are","in","the","comments","of","the","config","file",".","model","(","nn.Module",")",":","parent","MViT","module","this","module","forms","a","reversible","encoder","in","."]
code: def __init__(self, config, model):
      """
      The `__init__` method of any subclass should also contain these
          arguments.
      Args:
          cfg (CfgNode): model building configs, details are in the
              comments of the config file.
          model (nn.Module): parent MViT module this module forms
              a reversible encoder in.
      """
  
      super().__init__()
      self.cfg = config
  
      embed_dim = self.cfg.MVIT.EMBED_DIM
      depth = self.cfg.MVIT.DEPTH
      num_heads = self.cfg.MVIT.NUM_HEADS
      mlp_ratio = self.cfg.MVIT.MLP_RATIO
      qkv_bias = self.cfg.MVIT.QKV_BIAS
  
      drop_path_rate = self.cfg.MVIT.DROPPATH_RATE
      self.dropout = config.MVIT.DROPOUT_RATE
      self.pre_q_fusion = self.cfg.MVIT.REV.PRE_Q_FUSION
      dpr = [
          x.item() for x in torch.linspace(0, drop_path_rate, depth)
      ]  # stochastic depth decay rule
  
      input_size = model.patch_dims
  
      self.layers = nn.ModuleList([])
      self.no_custom_backward = False
  
      if self.cfg.MVIT.NORM == "layernorm":
          norm_layer = partial(nn.LayerNorm, eps=1e-6)
      else:
          raise NotImplementedError("Only supports layernorm.")
  
      dim_mul, head_mul = torch.ones(depth + 1), torch.ones(depth + 1)
      for i in range(len(self.cfg.MVIT.DIM_MUL)):
          dim_mul[self.cfg.MVIT.DIM_MUL[i][0]] = self.cfg.MVIT.DIM_MUL[
              i
          ][1]
      for i in range(len(self.cfg.MVIT.HEAD_MUL)):
          head_mul[
              self.cfg.MVIT.HEAD_MUL[i][0]
          ] = self.cfg.MVIT.HEAD_MUL[i][1]
  
      pool_q = model.pool_q
      pool_kv = model.pool_kv
      stride_q = model.stride_q
      stride_kv = model.stride_kv
  
      for i in range(depth):
          num_heads = round_width(num_heads, head_mul[i])
  
          # Upsampling inside the MHPA, input to the Q-pooling block is lower C dimension
          # This localizes the feature changes in a single block, making more computation reversible.
          embed_dim = round_width(
              embed_dim,
              dim_mul[i - 1] if i > 0 else 1.0,
              divisor=num_heads,
          )
          dim_out = round_width(
              embed_dim,
              dim_mul[i],
              divisor=round_width(num_heads, head_mul[i + 1]),
          )
  
          if i in self.cfg.MVIT.REV.BUFFER_LAYERS:
              layer_type = StageTransitionBlock
              input_mult = 2 if "concat" in self.pre_q_fusion else 1
          else:
              layer_type = ReversibleBlock
              input_mult = 1
  
          dimout_correction = (
              2
              if (input_mult == 2 and "concat" in self.pre_q_fusion)
              else 1
          )
  
          self.layers.append(
              layer_type(
                  dim=embed_dim
                  * input_mult,  # added only for concat fusion before Qpooling layers
                  input_size=input_size,
                  dim_out=dim_out * input_mult // dimout_correction,
                  num_heads=num_heads,
                  cfg=self.cfg,
                  mlp_ratio=mlp_ratio,
                  qkv_bias=qkv_bias,
                  drop_path=dpr[i],
                  norm_layer=norm_layer,
                  kernel_q=pool_q[i] if len(pool_q) > i else [],
                  kernel_kv=pool_kv[i] if len(pool_kv) > i else [],
                  stride_q=stride_q[i] if len(stride_q) > i else [],
                  stride_kv=stride_kv[i] if len(stride_kv) > i else [],
                  layer_id=i,
                  pre_q_fusion=self.pre_q_fusion,
              )
          )
          # F is the attention block
          self.layers[-1].F.thw = input_size
  
          if len(stride_q[i]) > 0:
              input_size = [
                  size // stride
                  for size, stride in zip(input_size, stride_q[i])
              ]
  
      embed_dim = dim_out
  
imports: import sys
  from functools import partial
  import torch
  from torch import nn
  from torch.autograd import Function
  from slowfast.models.attention import MultiScaleAttention, attention_pool
  from slowfast.models.common import Mlp, TwoStreamFusion, drop_path
  from slowfast.models.utils import round_width
label_desc: Use image node_id 1 to create a new ReversibleMViT object from inherited base classes: nn with example: obj = ReversibleMViT(config, model)