code: def compute_identity(edge_index, n, k):
      id, value = norm(edge_index, n)
      adj_sparse = torch.sparse.FloatTensor(
          id, value, torch.Size([n, n])
      )
      adj = adj_sparse.to_dense()
      diag_all = [torch.diag(adj)]
      adj_power = adj
      for i in range(1, k):
          adj_power = adj_power @ adj
          diag_all.append(torch.diag(adj_power))
      diag_all = torch.stack(diag_all, dim=1)
      return diag_all
  
imports: import torch
  from torch_geometric.utils import add_remaining_self_loops
  from torch_scatter import scatter_add
label_desc: Use image node_id 2 for calling a global function with example usage: compute_identity(edge_index, n, k) and returns: diag_all