cunumpy 0.7.0__tar.gz → 0.7.1__tar.gz

{cunumpy-0.7.0/src/cunumpy.egg-info → cunumpy-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cunumpy
-Version: 0.7.0
+Version: 0.7.1
 Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
 Author: Max
 License: MIT License

{cunumpy-0.7.0 → cunumpy-0.7.1}/pyproject.toml

@@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ]
 
 [project]
 name = "cunumpy"
-version = "0.7.0"
+version = "0.7.1"
 description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`."
 readme = "README.md"
 keywords = [ "python" ]

{cunumpy-0.7.0 → cunumpy-0.7.1}/src/cunumpy/_algorithms.py

@@ -380,6 +380,8 @@ def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]:
     radix sort; on NumPy, integer keys of more than 4096 entries get an LSD
     radix sort too (one stable ``uint16`` argsort per 16-bit digit the key range
     needs), about ten times faster than NumPy's stable sort of 64-bit integers.
+    On NumPy, a 2-D array sorted along its last axis (component-major markers) is
+    gathered row by row, about twice as fast as ``take`` along that axis.
 
     Parameters
     ----------
@@ -437,12 +439,27 @@ def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]:
         keys[order],
         order,
         *(
-            xpm.take(array, order, axis=array_axis)
+            _take_along_axis(xpm, array, order, array_axis)
             for array, array_axis in zip(arrays, axes, strict=True)
         ),
     )
 
 
+def _take_along_axis(xpm: Any, array: Any, order: Any, axis: int) -> Any:
+    """Return `array` taken at `order` along `axis`, as a new C-contiguous array.
+
+    A component-major ``(ncomp, n)`` NumPy array is gathered one contiguous row at a
+    time, which is about twice as fast as ``take`` along its last axis; everything
+    else, and CuPy, goes through ``take``.
+    """
+    if xpm is np and array.ndim == 2 and axis == 1:
+        gathered = np.empty((array.shape[0], order.shape[0]), dtype=array.dtype)
+        for row in range(array.shape[0]):
+            gathered[row] = array[row][order]
+        return gathered
+    return xpm.take(array, order, axis=axis)
+
+
 def compact_by_mask(mask: Any, *arrays: Any, axis: int = 0) -> int:
     """Move the entries where `mask` is True to the front of every array, in place.
 

{cunumpy-0.7.0 → cunumpy-0.7.1/src/cunumpy.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cunumpy
-Version: 0.7.0
+Version: 0.7.1
 Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
 Author: Max
 License: MIT License