code: def sort_key_val(keys, values, dimension=-1):
      idxs = list(np.ix_(*[np.arange(d) for d in keys.shape]))
      idxs[dimension] = np.argsort(keys, axis=dimension)
      return keys[tuple(idxs)], values[tuple(idxs)]
  
imports: import builtins
  import collections
  import itertools
  import numpy
  import opt_einsum
  import scipy.special
  from jax._src import dtypes
  from jax._src import util
label_desc: Use image node_id 35 for calling a global function with example usage: sort_key_val(keys, values, dimension) and returns: keys,  values