fix (JAX backend)(sorting.py): fixing the sort function to handle `…

…jax.__version__ > 0.4.28` wherein the `kind` argument has been deprecated.

ivy/functional/backends/jax/sorting.py

@@ -1,7 +1,8 @@
 # global
+import jax
 import jax.numpy as jnp
 from typing import Optional, Literal, Union, List
-
+from packaging.version import parse as _parse
 # local
 import ivy
 from ivy.func_wrapper import with_unsupported_dtypes
@@ -31,7 +32,11 @@ def sort(
     out: Optional[JaxArray] = None,
 ) -> JaxArray:
     kind = "stable" if stable else "quicksort"
-    ret = jnp.asarray(jnp.sort(x, axis=axis, kind=kind))
+    if _parse(jax.__version__) > _parse("0.4.28"):
+        # kind argument has been removed in jax 0.4.28 
+        ret = jnp.asarray(jnp.sort(x, axis=axis, stable=stable))
+    else:
+        ret = jnp.asarray(jnp.sort(x, axis=axis, kind=kind))
     if descending:
         ret = jnp.asarray(jnp.flip(ret, axis=axis))
     return ret