code: def test_edge_pooling():
      x = torch.tensor(
          [[0.0], [1.0], [2.0], [3.0], [4.0], [5.0], [-1.0]]
      )
      edge_index = torch.tensor(
          [
              [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5, 6],
              [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2, 5, 4, 0],
          ]
      )
      batch = torch.tensor([0, 0, 0, 0, 1, 1, 0])
  
      op = EdgePooling(in_channels=1)
      assert str(op) == "EdgePooling(1)"
  
      # Setting parameters fixed so we can test the expected outcome:
      op.lin.weight.data.fill_(1.0)
      op.lin.bias.data.fill_(0.0)
  
      # Test pooling:
      new_x, new_edge_index, new_batch, unpool_info = op(
          x, edge_index, batch
      )
  
      assert new_x.size(0) == new_batch.size(0) == 4
      assert new_edge_index.tolist() == [
          [0, 1, 1, 2, 2, 3],
          [0, 1, 2, 1, 2, 2],
      ]
      assert new_batch.tolist() == [1, 0, 0, 0]
  
      if is_full_test():
          jit = torch.jit.script(op)
          out = jit(x, edge_index, batch)
          assert torch.allclose(new_x, out[0])
          assert torch.equal(new_edge_index, out[1])
          assert torch.equal(new_batch, out[2])
  
      # Test unpooling:
      out = op.unpool(new_x, unpool_info)
      assert out[0].size() == x.size()
      assert out[0].tolist() == [[1], [1], [5], [5], [9], [9], [-1]]
      assert torch.equal(out[1], edge_index)
      assert torch.equal(out[2], batch)
  
      if is_full_test():
          jit = torch.jit.export(op)
          out = jit.unpool(new_x, unpool_info)
          assert out[0].size() == x.size()
          assert out[0].tolist() == [[1], [1], [5], [5], [9], [9], [-1]]
          assert torch.equal(out[1], edge_index)
          assert torch.equal(out[2], batch)
  
      # Test edge cases.
      x = torch.tensor([[0.0], [1.0], [2.0], [3.0], [4.0], [5.0]])
      edge_index = torch.tensor(
          [
              [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5],
              [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2, 5, 4],
          ]
      )
      batch = torch.tensor([0, 0, 0, 0, 1, 1])
      new_x, new_edge_index, new_batch, _ = op(x, edge_index, batch)
  
      assert new_x.size(0) == new_batch.size(0) == 3
      assert new_batch.tolist() == [1, 0, 0]
      assert new_edge_index.tolist() == [
          [0, 1, 1, 2, 2],
          [0, 1, 2, 1, 2],
      ]
  
imports: import torch
  from torch_geometric.nn import EdgePooling
  from torch_geometric.testing import is_full_test
  from torch_geometric.utils import scatter
label_desc: Use image node_id 4 for calling a global function with example usage: test_edge_pooling() without return types