class_docstr: A transformer layer among all nodes in a graph.
class_docstr_tok: ["A","transformer","layer","among","all","nodes","in","a","graph","."]
code: def __iter_idxs_per_graph(self, node_to_graph_idx):
      with torch.no_grad():
          num_nodes_per_graph = scatter_sum(
              src=torch.ones_like(node_to_graph_idx, dtype=torch.int64),
              index=node_to_graph_idx,
          )  # [num_nodes]
  
          node_offset = 0
          for num_nodes in num_nodes_per_graph:
              for start_idx in range(
                  0, num_nodes, self.__max_num_nodes
              ):
                  end_idx = min(
                      start_idx + self.__max_num_nodes, num_nodes
                  )
                  yield torch.arange(
                      start=start_idx,
                      end=end_idx,
                      dtype=torch.int64,
                      device=num_nodes_per_graph.device,
                  ) + node_offset
              node_offset += num_nodes
  
file_path: ptgnn/ptgnn/neuralmodels/gnn/messagepassing/selfattmessagepassing.py
imports: import torch
  from torch import nn
  from torch_scatter import scatter_sum
  from typing import Dict, List, Tuple
  from ptgnn.neuralmodels.gnn.messagepassing.abstractmessagepassing import AbstractMessagePassingLayer
label_desc: Use image node_id 2 for calling the MultiHeadSelfAttentionMessagePassing obj's underlying member method code with example usage: obj.__iter_idxs_per_graph(node_to_graph_idx) without return types