code: def test_sort_aggregation():
      N_1, N_2 = 4, 6
      x = torch.randn(N_1 + N_2, 4)
      index = torch.tensor(
          [0 for _ in range(N_1)] + [1 for _ in range(N_2)]
      )
  
      aggr = SortAggregation(k=5)
      assert str(aggr) == "SortAggregation(k=5)"
  
      out = aggr(x, index)
      assert out.size() == (2, 5 * 4)
  
      out_dim = out = aggr(x, index, dim=0)
      assert torch.allclose(out_dim, out)
  
      out = out.view(2, 5, 4)
  
      # First graph output has been filled up with zeros.
      assert out[0, -1].tolist() == [0, 0, 0, 0]
  
      # Nodes are sorted.
      assert torch.equal(out[0, :4, -1].argsort(), 3 - torch.arange(4))
      assert torch.equal(out[1, :, -1].argsort(), 4 - torch.arange(5))
  
imports: import torch
  from torch_geometric.nn.aggr import SortAggregation
label_desc: Use image node_id 1 for calling a global function with example usage: test_sort_aggregation() without return types