code: def test_avg_pool_x():
      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],
          ]
      )
      batch = torch.tensor([0, 0, 0, 0, 1, 1])
  
      out = avg_pool_x(cluster, x, batch)
      assert out[0].tolist() == [[3, 4], [5, 6], [10, 11]]
      assert out[1].tolist() == [0, 0, 1]
  
      if is_full_test():
          jit = torch.jit.script(avg_pool_x)
          out = jit(cluster, x, batch)
          assert out[0].tolist() == [[3, 4], [5, 6], [10, 11]]
          assert out[1].tolist() == [0, 0, 1]
  
      out, _ = avg_pool_x(cluster, x, batch, size=2)
      assert out.tolist() == [[3, 4], [5, 6], [10, 11], [0, 0]]
  
      batch_size = int(batch.max().item()) + 1
      out2, _ = avg_pool_x(
          cluster, x, batch, batch_size=batch_size, size=2
      )
      assert torch.equal(out, out2)
  
      if is_full_test():
          jit = torch.jit.script(avg_pool_x)
          out, _ = jit(cluster, x, batch, size=2)
          assert out.tolist() == [[3, 4], [5, 6], [10, 11], [0, 0]]
  
          out2, _ = jit(
              cluster, x, batch, batch_size=batch_size, size=2
          )
          assert torch.equal(out, out2)
  
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 1 for calling a global function with example usage: test_avg_pool_x() without return types