code: def ref_sort_bw(*args, **kwargs):
      if len(args) > 6:
          args = args[0:2] + args[3:]
      x, g, axis, reverse, with_index, only_index = args
      # Move sort axis to last and reshape into two dimensions
      xx = np.moveaxis(x, axis, -1).reshape(-1, x.shape[axis])
      gg = np.moveaxis(g, axis, -1).reshape(-1, x.shape[axis])
      # Argsort the last dimension in normal or reverse order
      ii = (
          np.argsort(xx) if not reverse else np.flip(np.argsort(xx), -1)
      )
      # Sort gradient vector over outer dimension
      for k in range(ii.shape[0]):
          gg[k] = gg[k, ii[k]]
      # Restore full shape but with sort axis at last
      gg = gg.reshape(np.moveaxis(g, axis, -1).shape)
      # Move the sort axis back to original position
      gg = np.moveaxis(gg, -1, axis)
      # The function tester expects a flattened gradient.
      return gg.flatten()
  
file_path: nnabla/python/test/function/test_sort.py
imports: import pytest
  import numpy
  import nnabla
  import nnabla.functions
  from nbla_test_utils import list_context, function_tester
label_desc: Use image node_id 2 for calling a global function with example usage: ref_sort_bw() and returns: gg