Remove deprecated kind argument from jnp.sort and jnp.argsort.

PiperOrigin-RevId: 631429900

CHANGELOG.md

@@ -69,6 +69,8 @@ Remember to align the itemized text with the first line of an item within a list
     {func}`~jax.numpy.apply_over_axes`, {func}`~jax.numpy.inner`,
     {func}`~jax.numpy.outer`, {func}`~jax.numpy.cross`,
     {func}`~jax.numpy.kron`, and {func}`~jax.numpy.lexsort`.
+  * The ``kind`` argument to {func}`jax.numpy.sort` and {func}`jax.numpy.argsort`
+    is now removed. Use `stable=True` or `stable=False` instead.
 
 * Bug fixes
   * {func}`jax.numpy.astype` will now always return a copy when `copy=True`.

jax/_src/numpy/lax_numpy.py

@@ -4502,30 +4502,29 @@ def _nanargmin(a, axis: int | None = None, keepdims : bool = False):
 
 
 @util.implements(np.sort, extra_params="""
-kind : deprecated; specify sort algorithm using stable=True or stable=False
-order : not supported
 stable : bool, default=True
     Specify whether to use a stable sort.
 descending : bool, default=False
     Specify whether to do a descending sort.
-    """)
+kind : deprecated; specify sort algorithm using stable=True or stable=False
+order : not supported
+""")
 @partial(jit, static_argnames=('axis', 'kind', 'order', 'stable', 'descending'))
 def sort(
     a: ArrayLike,
     axis: int | None = -1,
-    kind: str | None = None,
-    order: None = None, *,
+    *,
+    kind: None = None,
+    order: None = None,
     stable: bool = True,
     descending: bool = False,
 ) -> Array:
   util.check_arraylike("sort", a)
   if kind is not None:
-    # Deprecated 2024-01-05
-    warnings.warn("The 'kind' argument to sort has no effect, and is deprecated. "
-                  "Use stable=True or stable=False to specify sort stability.",
-                  category=DeprecationWarning, stacklevel=2)
+    raise TypeError("'kind' argument to sort is not supported. Use"
+                    " stable=True or stable=False to specify sort stability.")
   if order is not None:
-    raise ValueError("'order' argument to sort is not supported.")
+    raise TypeError("'order' argument to sort is not supported.")
   if axis is None:
     arr = ravel(a)
     axis = 0
@@ -4562,31 +4561,30 @@ def lexsort(keys: Array | np.ndarray | Sequence[ArrayLike], axis: int = -1) -> A
 
 
 @util.implements(np.argsort, extra_params="""
-kind : deprecated; specify sort algorithm using stable=True or stable=False
-order : not supported
 stable : bool, default=True
     Specify whether to use a stable sort.
 descending : bool, default=False
     Specify whether to do a descending sort.
+kind : deprecated; specify sort algorithm using stable=True or stable=False
+order : not supported
     """)
 @partial(jit, static_argnames=('axis', 'kind', 'order', 'stable', 'descending'))
 def argsort(
     a: ArrayLike,
     axis: int | None = -1,
-    kind: str | None = None,
+    *,
+    kind: None = None,
     order: None = None,
-    *, stable: bool = True,
+    stable: bool = True,
     descending: bool = False,
 ) -> Array:
   util.check_arraylike("argsort", a)
   arr = asarray(a)
   if kind is not None:
-    # Deprecated 2024-01-05
-    warnings.warn("The 'kind' argument to argsort has no effect, and is deprecated. "
-                  "Use stable=True or stable=False to specify sort stability.",
-                  category=DeprecationWarning, stacklevel=2)
+    raise TypeError("'kind' argument to argsort is not supported. Use"
+                    " stable=True or stable=False to specify sort stability.")
   if order is not None:
-    raise ValueError("'order' argument to argsort is not supported.")
+    raise TypeError("'order' argument to argsort is not supported.")
   if axis is None:
     arr = ravel(arr)
     axis = 0

jax/numpy/__init__.pyi

@@ -83,11 +83,11 @@ def argpartition(a: ArrayLike, kth: int, axis: int = ...) -> Array: ...
 def argsort(
     a: ArrayLike,
     axis: Optional[int] = ...,
-    kind: str | None = ...,
-    order: None = ...,
     *,
     stable: builtins.bool = ...,
     descending: builtins.bool = ...,
+    kind: str | None = ...,
+    order: None = ...,
 ) -> Array: ...
 def argwhere(
     a: ArrayLike,
@@ -737,11 +737,11 @@ sometrue = any
 def sort(
     a: ArrayLike,
     axis: Optional[int] = ...,
-    kind: str | None = ...,
-    order: None = ...,
     *,
     stable: builtins.bool = ...,
     descending: builtins.bool = ...,
+    kind: str | None = ...,
+    order: None = ...,
 ) -> Array: ...
 def sort_complex(a: ArrayLike) -> Array: ...
 def split(