code: def mask_edges(
      input_graph: BatchGraphData,
      mask_rate: float,
      masked_edge_indices=None,
  ):
      """
      Mask edges in a BatchGraphData object.
  
      This is separate from the mask_nodes function because we want to be able to mask edges without masking any nodes.
  
      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_edge_indices : list, optional
          If None, then randomly sample num_edges * mask_rate + 1 number of edge indices. Otherwise should correspond to the 1st direction of an edge pair. ie all indices should be an even number
  
      Returns
      -------
      data: dc.feat.BatchGraphData
          Creates new attributes in the original  object:
          - data.mask_edge_idx: indices of masked edges
          - data.mask_edge_labels: corresponding ground truth edge feature for each masked edge
          - data.edge_attr: modified in place: the edge features (both directions) that correspond to the masked edges have the masked edge feature
      """
      data = copy.deepcopy(input_graph)
  
      if masked_edge_indices is None:
          # sample x distinct edges to be masked, based on mask rate. But
          # will sample at least 1 edge
          num_edges = int(
              data.edge_index.size()[1] / 2  # type: ignore
          )  # num unique edges
          sample_size = int(num_edges * mask_rate + 1)
          # during sampling, we only pick the 1st direction of a particular
          # edge pair
          masked_edge_indices = [
              2 * i
              for i in random.sample(range(num_edges), sample_size)
          ]
  
      data.masked_edge_idx = torch.tensor(  # type: ignore
          np.array(masked_edge_indices)
      )
  
      # create ground truth edge features for the edges that correspond to
      # the masked indices
      mask_edge_labels_list = []
      for idx in masked_edge_indices:
          mask_edge_labels_list.append(  # yapf: disable
              data.edge_features[idx].view(1, -1)  # type: ignore
          )
      data.mask_edge_label = torch.cat(  # type: ignore
          mask_edge_labels_list, dim=0
      )
  
      # created new masked edge_attr, where both directions of the masked
      # edges have masked edge type. For message passing in gcn
  
      # append the 2nd direction of the masked edges
      all_masked_edge_indices = masked_edge_indices + [
          i + 1 for i in masked_edge_indices
      ]
      num_edge_feat = data.edge_features.size()[1]  # type: ignore
      for idx in all_masked_edge_indices:
          data.edge_features[idx] = torch.zeros(  # type: ignore
              (1, num_edge_feat)
          )
      # zeros are meant to represent the masked features. This is distinct from the
      # original implementation, where the masked features are represented by 0s and
      # an additional mask feature
      # link to source: https://github.com/snap-stanford/pretrain-gnns/blob/08f126ac13623e551a396dd5e511d766f9d4f8ff/bio/util.py#L101
  
      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 3 for calling a global function with example usage: mask_edges(input_graph, mask_rate, masked_edge_indices) and returns: data