cunumpy 0.7.1__tar.gz → 0.7.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cunumpy
-Version: 0.7.1
+Version: 0.7.2
 Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
 Author: Max
 License: MIT License
@@ -40,6 +40,7 @@ License-File: LICENSE.txt
 Requires-Dist: array-api-compat
 Requires-Dist: maybempi>=0.1.2
 Requires-Dist: numpy
+Requires-Dist: scope-profiler>=0.7.5
 Provides-Extra: dev
 Requires-Dist: cunumpy[docs,test-compiled]; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
@@ -50,6 +51,12 @@ Requires-Dist: griffe; extra == "docs"
 Requires-Dist: ipykernel; extra == "docs"
 Provides-Extra: metal
 Requires-Dist: mlx; (sys_platform == "darwin" and platform_machine == "arm64") and extra == "metal"
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
+Provides-Extra: mlx
+Requires-Dist: mlx; (sys_platform == "darwin" and platform_machine == "arm64") and extra == "mlx"
+Provides-Extra: torch
+Requires-Dist: torch; extra == "torch"
 Provides-Extra: test
 Requires-Dist: coverage; extra == "test"
 Requires-Dist: pytest; extra == "test"

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

@@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ]
 
 [project]
 name = "cunumpy"
-version = "0.7.1"
+version = "0.7.2"
 description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`."
 readme = "README.md"
 keywords = [ "python" ]
@@ -25,6 +25,7 @@ dependencies = [
   "array-api-compat",
   "maybempi>=0.1.2",
   "numpy",
+  "scope-profiler>=0.7.5",
 ]
 
 optional-dependencies.dev = [
@@ -37,6 +38,9 @@ optional-dependencies.dev = [
 # the shared docs workflow registers an ipykernel even without notebooks
 optional-dependencies.docs = [ "griffe", "ipykernel" ]
 optional-dependencies.metal = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ]
+optional-dependencies.jax = [ "jax" ]
+optional-dependencies.mlx = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ]
+optional-dependencies.torch = [ "torch" ]
 optional-dependencies.test = [ "coverage", "pytest", "scipy" ]
 optional-dependencies.test-compiled = [ "cunumpy[test]", "pyccel" ]
 urls."Source" = "https://github.com/max-models/cunumpy"
@@ -47,6 +51,12 @@ where = [ "src" ]
 [tool.setuptools.package-data]
 cunumpy = [ "py.typed", "*.pyi", "LLM_GUIDE.md", "cuda/include/cunumpy/*.cuh" ]
 
+# PRs to devel test against the devel branches of the max-models packages; PRs to main
+# and the release use PyPI only.
+[tool.uv.sources]
+maybempi = { git = "https://github.com/max-models/maybempi", branch = "devel" }
+scope-profiler = { git = "https://github.com/max-models/scope-profiler", branch = "devel" }
+
 [tool.ruff]
 # agent worktrees and scratch scripts are local copies, not part of the project
 extend-exclude = [ ".claude" ]

{cunumpy-0.7.1 → cunumpy-0.7.2}/src/cunumpy/__init__.py

@@ -15,11 +15,14 @@ active backend is NumPy or CuPy (prefer :func:`get_backend` for the name).
 import re as _re
 from importlib.metadata import PackageNotFoundError, version
 
+import scope_profiler as _sp
+
 from cunumpy import (
     algorithms,
     arguments,
     cuda,
     integrate,
+    interop,
     kernels,
     memory,
     mpi,
@@ -30,6 +33,7 @@ from cunumpy import (
     xp,
 )
 from cunumpy._host import evaluate_on_host, host_call, setup_on_host
+from cunumpy._profile_tags import _PROFILE_TAGS
 from cunumpy._scipy_backend import scipy
 from cunumpy._staging import to_host_async
 from cunumpy.xp import (
@@ -54,6 +58,10 @@ from cunumpy.xp import (
     use_backend,
 )
 
+# cunumpy's scope-profiler regions are off unless the application enables the
+# "cunumpy" tag
+_sp.register_opt_in_tag(*_PROFILE_TAGS)
+
 try:
     __version__ = version("cunumpy")
     """The installed version of cunumpy; ``"0.0.0+unknown"`` without package metadata."""
@@ -116,6 +124,7 @@ __all__ = [
     "get_backend",
     "host_call",
     "integrate",
+    "interop",
     "is_cpu",
     "is_gpu",
     "kernels",

{cunumpy-0.7.1 → cunumpy-0.7.2}/src/cunumpy/__init__.pyi

@@ -12,6 +12,7 @@ from cunumpy import algorithms as algorithms
 from cunumpy import arguments as arguments
 from cunumpy import cuda as cuda
 from cunumpy import integrate as integrate
+from cunumpy import interop as interop
 from cunumpy import kernels as kernels
 from cunumpy import memory as memory
 from cunumpy import mpi as mpi

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

@@ -7,8 +7,10 @@ import operator
 from typing import Any
 
 import array_api_compat.numpy as np
+import scope_profiler as sp
 
 from cunumpy import _fake_cupy
+from cunumpy._profile_tags import _PROFILE_TAGS
 from cunumpy.xp import assert_same_backend, get_array_backend, get_array_module
 
 
@@ -202,6 +204,7 @@ class SegmentPlan:
         """Number of output segments, fixed when the plan is created."""
         return self._n_segments
 
+    @sp.profile("cunumpy.SegmentPlan.sum", tags=_PROFILE_TAGS)
     def sum(self, values: Any, *, out: Any = None) -> Any:
         """Sum the rows of `values` per segment.
 
@@ -371,6 +374,22 @@ def _radix_argsort(keys: np.ndarray) -> np.ndarray:
     return order.astype(np.int64, copy=False)
 
 
+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)
+
+
+@sp.profile("cunumpy.sort_by_key", tags=_PROFILE_TAGS)
 def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]:
     """Sort `keys` and reorder every array the same way, in one stable argsort.
 
@@ -445,21 +464,7 @@ def sort_by_key(keys: Any, *arrays: Any, axis: int = 0) -> tuple[Any, ...]:
     )
 
 
-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)
-
-
+@sp.profile("cunumpy.compact_by_mask", tags=_PROFILE_TAGS)
 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.