code: def masked_edge_index(edge_index: Adj, edge_mask: Tensor) -> Adj:
      if isinstance(edge_index, Tensor):
          return edge_index[:, edge_mask]
      return torch_sparse.masked_select_nnz(
          edge_index, edge_mask, layout="coo"
      )
  
imports: from typing import Optional, Tuple, Union
  import torch
  from torch import Tensor
  from torch.nn import Parameter
  from torch.nn import Parameter
  import torch_geometric.backend
  import torch_geometric.typing
  from torch_geometric import is_compiling
  from torch_geometric.nn.conv import MessagePassing
  from torch_geometric.nn.inits import glorot, zeros
  from torch_geometric.typing import Adj, OptTensor, SparseTensor, pyg_lib, torch_sparse
  from torch_geometric.utils import index_sort, one_hot, scatter, spmm
  from torch_geometric.utils.sparse import index2ptr
label_desc: Use image node_id 1 for calling a global function with example usage: masked_edge_index(edge_index, edge_mask) and returns: torch_sparse, edge_index