import torch


def find_neighbor_edges(g, node_id):
    """Given a node with its graph, return all the edges connected to that node."""
    predecessors = g.predecessors(node_id)
    successors = g.successors(node_id)
    predecessors_edges = g.edge_ids(predecessors, torch.full(predecessors.shape, node_id, dtype=torch.int))
    successors_edges = g.edge_ids(torch.full(successors.shape, node_id, dtype=torch.int), successors)
    return torch.cat((predecessors_edges, successors_edges))