code: def test_mem_pool():
      mpool1 = MemPooling(4, 8, heads=3, num_clusters=2)
      assert str(mpool1) == "MemPooling(4, 8, heads=3, num_clusters=2)"
      mpool2 = MemPooling(8, 4, heads=2, num_clusters=1)
  
      x = torch.randn(17, 4)
      batch = torch.tensor(
          [0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 4, 4, 4]
      )
      _, mask = to_dense_batch(x, batch)
  
      out1, S = mpool1(x, batch)
      loss = MemPooling.kl_loss(S)
      with torch.autograd.set_detect_anomaly(True):
          loss.backward()
      out2, _ = mpool2(out1)
  
      assert out1.size() == (5, 2, 8)
      assert out2.size() == (5, 1, 4)
      assert S[~mask].sum() == 0
      assert round(S[mask].sum().item()) == x.size(0)
      assert float(loss) > 0
  
label_desc: Use image node_id 1 for calling a global function with example usage: test_mem_pool() without return types