code: def compute_rollout_attention(all_layer_matrices, start_layer=0):
      # adding residual consideration
      num_tokens = all_layer_matrices[0].shape[1]
      batch_size = all_layer_matrices[0].shape[0]
      eye = (
          torch.eye(num_tokens)
          .expand(batch_size, num_tokens, num_tokens)
          .to(all_layer_matrices[0].device)
      )
      all_layer_matrices = [
          all_layer_matrices[i] + eye
          for i in range(len(all_layer_matrices))
      ]
      # all_layer_matrices = [all_layer_matrices[i] / all_layer_matrices[i].sum(dim=-1, keepdim=True)
      #                       for i in range(len(all_layer_matrices))]
      joint_attention = all_layer_matrices[start_layer]
      for i in range(start_layer + 1, len(all_layer_matrices)):
          joint_attention = all_layer_matrices[i].bmm(joint_attention)
      return joint_attention
  
file_path: towhee/towhee/models/vit/vit.py
imports: import torch
  from torch import nn
  from torch.utils import model_zoo
  from functools import partial
  from collections import OrderedDict
  from towhee.models.utils.init_vit_weights import init_vit_weights
  from towhee.models.layers.patch_embed2d import PatchEmbed2D
  from .vit_utils import get_configs
  from .vit_block import Block
  from towhee.models.layers.layers_with_relprop import LayerNorm, GELU, Linear, IndexSelect, Add
label_desc: Use image node_id 1 for calling a global function with example usage: compute_rollout_attention(all_layer_matrices, start_layer) and returns: joint_attention