code: def sort_tensor(
      context: torch.Tensor,
      lens: torch.Tensor,
      dim: Optional[int] = 0,
      descending: Optional[bool] = True,
  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
      """Sorts elements in context by the dim lengths specified in lens
      Args:
          context:  source tensor, sorted by lens
          lens: lengths of elements of context along the dimension dim
          dim: Optional[int] : dimension to sort by
      Returns:
          context: tensor sorted by lens along dimension dim
          lens_sorted: lens tensor, sorted
          ids_sorted: reorder ids to be used to restore original order
  
      """
      lens_sorted, ids_sorted = torch.sort(lens, descending=descending)
      context = torch.index_select(context, dim, ids_sorted)
      return context, lens_sorted, ids_sorted
  
imports: from enum import Enum
  from typing import Optional, Tuple
  import librosa
  import matplotlib.pylab
  import numpy
  import torch
  from einops import rearrange
  from numba import jit, prange
  from nemo.collections.tts.torch.tts_data_types import DATA_STR2DATA_CLASS, MAIN_DATA_TYPES, WithLens
  from nemo.utils import logging
  from nemo.utils.decorators import deprecated
label_desc: Use image node_id 6 for calling a global function with example usage: sort_tensor(context, lens, dim, descending) and returns: context,  lens_sorted,  ids_sorted