MAINT: use copyto in ndarray.sort, nuke _ufuncs_impl._sort

Author: ev-br

torch_np/_funcs.py

@@ -3,7 +3,7 @@
 from . import _funcs_impl
 from ._normalizations import normalizer
 
-# _funcs_imple.py contains functions which mimic NumPy's eponimous equivalents,
+# _funcs_impl.py contains functions which mimic NumPy's eponimous equivalents,
 # and consume/return PyTorch tensors/dtypes.
 # They are also type annotated.
 # Pull these functions from _funcs_impl and decorate them with @normalizer, which

torch_np/_funcs_impl.py

@@ -1237,15 +1237,10 @@ def _sort_helper(tensor, axis, kind, order):
     return tensor, axis, stable
 
 
-def _sort(tensor, axis, kind, order):
-    # pure torch implementation, used below and in ndarray.sort
-    tensor, axis, stable = _sort_helper(tensor, axis, kind, order)
-    result = torch.sort(tensor, dim=axis, stable=stable)
-    return result.values
-
-
 def sort(a: ArrayLike, axis=-1, kind=None, order=None):
-    return _sort(a, axis, kind, order)
+    a, axis, stable = _sort_helper(a, axis, kind, order)
+    result = torch.sort(a, dim=axis, stable=stable)
+    return result.values
 
 
 def argsort(a: ArrayLike, axis=-1, kind=None, order=None):

torch_np/_ndarray.py

@@ -376,7 +376,7 @@ def resize(self, *new_shape, refcheck=False):
 
     def sort(self, axis=-1, kind=None, order=None):
         # ndarray.sort works in-place
-        self.tensor.copy_(_funcs_impl._sort(self.tensor, axis, kind, order))
+        _funcs.copyto(self, _funcs.sort(self, axis, kind, order))
 
     argsort = _funcs.argsort
     searchsorted = _funcs.searchsorted