code: def compute_rollout_attention(all_layer_matrices, start_layer=0):
      all_layer_matrices = align_scale(all_layer_matrices)
      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))
      ]
      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
  
imports: import math
  from functools import partial
  from typing import Tuple, List
  import numpy
  import torch
  from torch import nn
  from torch.nn.init import trunc_normal_
  from einops import rearrange
  from towhee.models.layers.layers_with_relprop import GELU, Linear, Dropout, Conv2d, Conv3d, Softmax, Sigmoid, LayerNorm, Einsum, Add, Clone, MaxPool3d, IndexSelect
label_desc: Use image node_id 4 for calling a global function with example usage: compute_rollout_attention(all_layer_matrices, start_layer) and returns: joint_attention