class_docstr: Base class for implementing rooted subgraph transformations.
class_docstr_tok: ["Base","class","for","implementing","rooted","subgraph","transformations","."]
code: def map(
      self,
      data: Data,
      n_mask: Tensor,
  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
      assert data.edge_index is not None
      num_nodes = data.num_nodes
      assert num_nodes is not None
  
      n_sub_batch, n_id = n_mask.nonzero().t()
      e_mask = (
          n_mask[:, data.edge_index[0]] & n_mask[:, data.edge_index[1]]
      )
      e_sub_batch, e_id = e_mask.nonzero().t()
  
      sub_edge_index = data.edge_index[:, e_id]
      arange = torch.arange(n_id.size(0), device=data.edge_index.device)
      node_map = data.edge_index.new_ones(num_nodes, num_nodes)
      node_map[n_sub_batch, n_id] = arange
      sub_edge_index += (arange * data.num_nodes)[e_sub_batch]
      sub_edge_index = node_map.view(-1)[sub_edge_index]
  
      return sub_edge_index, n_id, e_id, n_sub_batch, e_sub_batch
  
imports: import copy
  from abc import ABC, abstractmethod
  from typing import Any, Tuple
  import torch
  from torch import Tensor
  from torch_geometric.data import Data
  from torch_geometric.transforms import BaseTransform
  from torch_geometric.utils import to_torch_csc_tensor
label_desc: Use image node_id 2 for calling the RootedSubgraph obj's underlying member method code with example usage: obj.map(data, n_mask) and returns: sub_edge_index,  n_id,  e_id,  n_sub_batch,  e_sub_batch