class_docstr: Reversible Blocks for Reversible Vision Transformer and also
  for state-preserving blocks in Reversible MViT. See Section
  3.3.2 in paper for details.
class_docstr_tok: ["Reversible","Blocks","for","Reversible","Vision","Transformer","and","also","for","state-preserving","blocks","in","Reversible","MViT",".","See","Section","3.3.2","in","paper","for","details","."]
docstr: equation for activation recomputation:
  X_2 = Y_2 - G(Y_1), G = MLP
  X_1 = Y_1 - F(X_2), F = Attention
docstr_tok: ["equation","for","activation","recomputation",":","X_2","=","Y_2","-","G","(","Y_1",")",",","G","=","MLP","X_1","=","Y_1","-","F","(","X_2",")",",","F","=","Attention"]
code: def backward_pass(
      self,
      Y_1,
      Y_2,
      dY_1,
      dY_2,
  ):
      """
      equation for activation recomputation:
      X_2 = Y_2 - G(Y_1), G = MLP
      X_1 = Y_1 - F(X_2), F = Attention
      """
  
      # temporarily record intermediate activation for G
      # and use them for gradient calculcation of G
      with torch.enable_grad():
          Y_1.requires_grad = True
  
          torch.manual_seed(self.seeds["FFN"])
          g_Y_1 = self.G(Y_1)
  
          torch.manual_seed(self.seeds["droppath"])
          g_Y_1 = drop_path(
              g_Y_1,
              drop_prob=self.drop_path_rate,
              training=self.training,
          )
  
          g_Y_1.backward(dY_2, retain_graph=True)
  
      # activation recomputation is by design and not part of
      # the computation graph in forward pass.
      with torch.no_grad():
          X_2 = Y_2 - g_Y_1
          del g_Y_1
  
          dY_1 = dY_1 + Y_1.grad
          Y_1.grad = None
  
      # record F activations and calc gradients on F
      with torch.enable_grad():
          X_2.requires_grad = True
  
          torch.manual_seed(self.seeds["attn"])
          f_X_2 = self.F(X_2)
  
          torch.manual_seed(self.seeds["droppath"])
          f_X_2 = drop_path(
              f_X_2,
              drop_prob=self.drop_path_rate,
              training=self.training,
          )
  
          f_X_2.backward(dY_1, retain_graph=True)
  
      # propagate reverse computed acitvations at the start of
      # the previou block for backprop.s
      with torch.no_grad():
          X_1 = Y_1 - f_X_2
  
          del f_X_2, Y_1
          dY_2 = dY_2 + X_2.grad
  
          X_2.grad = None
          X_2 = X_2.detach()
  
      return X_1, X_2, dY_1, dY_2
  
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 4 for calling the ReversibleBlock obj's underlying member method code with example usage: obj.backward_pass(Y_1, Y_2, dY_1, dY_2) and returns: X_1,  X_2,  dY_1,  dY_2