code: def test_sort_aggregation_dim_size():
      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)"
  
      # expand batch output by 1
      out = aggr(x, index, dim_size=3)
      assert out.size() == (3, 5 * 4)
  
      out = out.view(3, 5, 4)
  
      # Both first and last graph outputs have been filled up with zeros.
      assert out[0, -1].tolist() == [0, 0, 0, 0]
      assert out[2, -1].tolist() == [0, 0, 0, 0]
  
imports: import torch
  from torch_geometric.nn.aggr import SortAggregation
label_desc: Use image node_id 3 for calling a global function with example usage: test_sort_aggregation_dim_size() without return types