code: def test_avg_pool():
      cluster = torch.tensor([0, 1, 0, 1, 2, 2])
      x = torch.tensor(
          [
              [1.0, 2.0],
              [3.0, 4.0],
              [5.0, 6.0],
              [7.0, 8.0],
              [9.0, 10.0],
              [11.0, 12.0],
          ]
      )
      pos = torch.tensor(
          [
              [0.0, 0.0],
              [1.0, 1.0],
              [2.0, 2.0],
              [3.0, 3.0],
              [4.0, 4.0],
              [5.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],
          ]
      )
      edge_attr = torch.ones(edge_index.size(1))
      batch = torch.tensor([0, 0, 0, 0, 1, 1])
  
      data = Batch(
          x=x,
          pos=pos,
          edge_index=edge_index,
          edge_attr=edge_attr,
          batch=batch,
      )
  
      data = avg_pool(cluster, data, transform=lambda x: x)
  
      assert data.x.tolist() == [[3, 4], [5, 6], [10, 11]]
      assert data.pos.tolist() == [[1, 1], [2, 2], [4.5, 4.5]]
      assert data.edge_index.tolist() == [[0, 1], [1, 0]]
      assert data.edge_attr.tolist() == [4, 4]
      assert data.batch.tolist() == [0, 0, 1]
  
imports: import torch
  from torch_geometric.data import Batch
  from torch_geometric.nn import avg_pool, avg_pool_neighbor_x, avg_pool_x
  from torch_geometric.testing import is_full_test
label_desc: Use image node_id 2 for calling a global function with example usage: test_avg_pool() without return types