class_docstr: A general GNN layer that supports edge features and edge sampling.
class_docstr_tok: ["A","general","GNN","layer","that","supports","edge","features","and","edge","sampling","."]
code: def forward(self, batch):
      edge_mask = torch.rand(batch.edge_index.shape[1]) < self.keep_edge
      edge_index = batch.edge_index[:, edge_mask]
      edge_feature = batch.edge_attr[edge_mask, :]
      batch.x = self.model(
          batch.x, edge_index, edge_feature=edge_feature
      )
      return batch
  
imports: import copy
  from dataclasses import dataclass, replace
  import torch
  import torch.nn.functional
  import torch_geometric
  import torch_geometric.graphgym.models.act
  import torch_geometric.graphgym.register
  from torch_geometric.graphgym.contrib.layer.generalconv import GeneralConvLayer, GeneralEdgeConvLayer
  from torch_geometric.graphgym.register import register_layer
  from torch_geometric.nn import Linear
label_desc: Use image node_id 2 for calling the GeneralSampleEdgeConv obj's underlying member method code with example usage: obj.forward(batch) and returns: batch