code: def test_sort_edge_index_jit():
      @torch.jit.script
      def wrapper1(edge_index: Tensor) -> Tensor:
          return sort_edge_index(edge_index)
  
      @torch.jit.script
      def wrapper2(
          edge_index: Tensor,
          edge_attr: Optional[Tensor],
      ) -> Tuple[Tensor, Optional[Tensor]]:
          return sort_edge_index(edge_index, edge_attr)
  
      @torch.jit.script
      def wrapper3(
          edge_index: Tensor,
          edge_attr: List[Tensor],
      ) -> Tuple[Tensor, List[Tensor]]:
          return sort_edge_index(edge_index, edge_attr)
  
      edge_index = torch.tensor([[2, 1, 1, 0], [1, 2, 0, 1]])
      edge_attr = torch.tensor([[1], [2], [3], [4]])
  
      out = wrapper1(edge_index)
      assert out.size() == edge_index.size()
  
      out = wrapper2(edge_index, None)
      assert out[0].size() == edge_index.size()
      assert out[1] is None
  
      out = wrapper2(edge_index, edge_attr)
      assert out[0].size() == edge_index.size()
      assert out[1].size() == edge_attr.size()
  
      out = wrapper3(edge_index, [edge_attr, edge_attr.view(-1)])
      assert out[0].size() == edge_index.size()
      assert len(out[1]) == 2
      assert out[1][0].size() == edge_attr.size()
      assert out[1][1].size() == edge_attr.view(-1).size()
  
imports: from typing import List, Optional, Tuple
  import torch
  from torch import Tensor
  from torch_geometric.utils import sort_edge_index
label_desc: Use image node_id 2 for calling a global function with example usage: test_sort_edge_index_jit() and returns: sort_edge_index, sort_edge_index, sort_edge_index