code: def test_lightgcn_link_prediction(
      embedding_dim, with_edge_weight, alpha
  ):
      num_nodes = 500
      num_edges = 400
      edge_index = torch.randint(0, num_nodes, (2, num_edges))
      edge_weight = torch.rand(num_edges) if with_edge_weight else None
      edge_label_index = torch.randint(0, num_nodes, (2, 100))
      edge_label = torch.randint(0, 2, (edge_label_index.size(1),))
  
      model = LightGCN(
          num_nodes, embedding_dim, num_layers=2, alpha=alpha
      )
      assert (
          str(model) == f"LightGCN(500, {embedding_dim}, num_layers=2)"
      )
  
      pred = model(edge_index, edge_label_index, edge_weight)
      assert pred.size() == (100,)
  
      loss = model.link_pred_loss(pred, edge_label)
      assert loss.dim() == 0 and loss > 0
  
      prob = model.predict_link(
          edge_index, edge_label_index, edge_weight, prob=True
      )
      assert prob.size() == (100,)
      assert prob.min() > 0 and prob.max() < 1
  
      prob = model.predict_link(
          edge_index, edge_label_index, edge_weight, prob=False
      )
      assert prob.size() == (100,)
      assert ((prob == 0) | (prob == 1)).sum() == 100
  
imports: import pytest
  import torch
  from torch_geometric.nn.models import LightGCN
label_desc: Use image node_id 2 for calling a global function with example usage: test_lightgcn_link_prediction(embedding_dim, with_edge_weight, alpha) without return types