code: def mask_nodes(
      input_graph: BatchGraphData,
      mask_rate,
      masked_node_indices=None,
      mask_edge=True,
  ):
      """
      Mask nodes and their connected edges in a BatchGraphData object.
  
      This function assumes that the first node feature is the atomic number, for example with the SNAPFeaturizer. It will set masked nodes' features to 0.
  
      Parameters
      ----------
      input_graph: dc.feat.BatchGraphData
          Assume that the edge ordering is the default PyTorch geometric ordering, where the two directions of a single edge occur in pairs.
          Eg. data.edge_index = tensor([[0, 1, 1, 2, 2, 3],
                                      [1, 0, 2, 1, 3, 2]])
      masked_node_indices: list, optional
          If None, then randomly samples num_nodes * mask rate number of node indices. Otherwise, a list of node indices that sets the nodes to be masked (for debugging only).
      mask_edge: bool, optional
          Will mask the edges connected to the masked nodes.
  
      Returns
      -------
      data: dc.feat.BatchGraphData
          Creates new attributes in the original data object:
          - data.mask_node_idx
          - data.mask_node_label
          - data.mask_edge_idx
          - data.mask_edge_label
  
      """
      data = copy.deepcopy(input_graph)
  
      if masked_node_indices is None:
          # sample x distinct nodes to be masked, based on mask rate. But
          # will sample at least 1 node
          num_nodes = data.node_features.size()[0]  # type: ignore
          sample_size = int(num_nodes * mask_rate + 1)
          masked_node_indices = random.sample(
              range(num_nodes), sample_size
          )
  
      # create mask node label by copying node feature of mask node
      mask_node_labels_list = []
      for node_idx in masked_node_indices:
          mask_node_labels_list.append(
              data.node_features[node_idx].view(1, -1)
          )
      data.mask_node_label = torch.cat(  # type: ignore
          mask_node_labels_list, dim=0
      )[:, 0].long()
      data.masked_node_indices = torch.tensor(masked_node_indices)  # type: ignore
  
      # modify the original node feature of the masked node
      num_node_feats = data.node_features.size()[1]  # type: ignore
      for node_idx in masked_node_indices:
          data.node_features[node_idx] = torch.zeros(
              (1, num_node_feats)
          )
      # zeros are meant to represent the masked features. This is distinct from the
      # original implementation, where the masked features are represented by the
      # the last feature token 119.
      # link to source: https://github.com/snap-stanford/pretrain-gnns/blob/08f126ac13623e551a396dd5e511d766f9d4f8ff/chem/util.py#L241
  
      if mask_edge:
          # create mask edge labels by copying edge features of edges that are connected to
          # mask nodes
          connected_edge_indices = []
          for edge_idx, (u, v) in enumerate(
              data.edge_index.cpu().numpy().T
          ):  # type: ignore
              for node_idx in masked_node_indices:
                  if (
                      node_idx in set((u, v))
                      and edge_idx not in connected_edge_indices
                  ):
                      connected_edge_indices.append(edge_idx)
  
          if len(connected_edge_indices) > 0:
              # create mask edge labels by copying edge features of the edges connected to
              # the mask nodes
              mask_edge_labels_list = []
              for edge_idx in connected_edge_indices[
                  ::2
              ]:  # because the
                  # edge ordering is such that two directions of a single
                  # edge occur in pairs, so to get the unique undirected
                  # edge indices, we take every 2nd edge index from list
                  mask_edge_labels_list.append(
                      data.edge_features[edge_idx].view(  # type: ignore
                          1, -1
                      )
                  )
  
              data.mask_edge_label = torch.cat(  # type: ignore
                  mask_edge_labels_list, dim=0
              )[
                  :, 0
              ].long()  # type: ignore
              # modify the original edge features of the edges connected to the mask nodes
              num_edge_feat = data.edge_features.size()[1]  # type: ignore
              for edge_idx in connected_edge_indices:
                  data.edge_features[edge_idx] = torch.zeros(  # type: ignore
                      (1, num_edge_feat)
                  )  # type: ignore
              # zeros are meant to represent the masked features. This is distinct from the
              # original implementation, where the masked features are represented by the
              # the last feature token 4.
              # link to source: https://github.com/snap-stanford/pretrain-gnns/blob/08f126ac13623e551a396dd5e511d766f9d4f8ff/chem/util.py#L268
  
              data.connected_edge_indices = torch.tensor(  # type: ignore
                  connected_edge_indices[::2]
              )
          else:
              data.mask_edge_label = torch.empty(  # type: ignore
                  (0, 2)
              ).to(torch.int64)
              data.connected_edge_indices = torch.tensor(  # type: ignore
                  connected_edge_indices
              ).to(
                  torch.int64
              )
  
      return data
  
imports: import random
  import copy
  import torch
  import numpy
  from torch_geometric.nn import GINEConv, GCNConv, GATConv, SAGEConv, global_add_pool, global_mean_pool, global_max_pool
  from torch_geometric.nn.aggr import AttentionalAggregation, Set2Set
  from torch_geometric.nn.inits import uniform
  import torch.nn
  from torch.functional import F
  from deepchem.data import Dataset
  from deepchem.models.losses import SoftmaxCrossEntropy, EdgePredictionLoss, GraphNodeMaskingLoss, GraphEdgeMaskingLoss, DeepGraphInfomaxLoss, GraphContextPredLoss
  from deepchem.models.torch_models import ModularTorchModel
  from deepchem.feat.graph_data import BatchGraphData, GraphData, shortest_path_length
  from typing import Iterable, List, Tuple
  from deepchem.metrics import to_one_hot
label_desc: Use image node_id 2 for calling a global function with example usage: mask_nodes(input_graph, mask_rate, masked_node_indices, mask_edge) and returns: data