diff --git a/pyproject.toml b/pyproject.toml index ac9a151..60f12bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,5 +56,8 @@ cunumpy = [ "py.typed", "*.pyi", "LLM_GUIDE.md", "cuda/include/cunumpy/*.cuh" ] # agent worktrees and scratch scripts are local copies, not part of the project extend-exclude = [ ".claude" ] # formatting: `ruff format`; import sorting: the isort rules of `ruff check` -lint.extend-select = [ "I" ] +lint.extend-select = [ "I", "TID252" ] lint.isort.known-first-party = [ "cunumpy" ] + +[tool.ruff.lint.flake8-tidy-imports] +ban-relative-imports = "all" diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index b0e9da4..24297ab 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -3,9 +3,9 @@ import warnings as _warnings from importlib.metadata import PackageNotFoundError, version -from . import algorithms, cuda, kernels, memory, mpi, petsc, profiling, rng, xp -from ._scipy_backend import scipy -from .xp import ( +from cunumpy import algorithms, cuda, kernels, memory, mpi, petsc, profiling, rng, xp +from cunumpy._scipy_backend import scipy +from cunumpy.xp import ( as_device_array, assert_same_backend, backend_info, diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index eadad06..ad45ff0 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -8,16 +8,16 @@ from typing import Any import numpy as np from numpy import * -from . import algorithms as algorithms -from . import cuda as cuda -from . import kernels as kernels -from . import memory as memory -from . import mpi as mpi -from . import petsc as petsc -from . import profiling as profiling -from . import rng as rng -from . import xp as xp -from ._scipy_backend import scipy as scipy +from cunumpy import algorithms as algorithms +from cunumpy import cuda as cuda +from cunumpy import kernels as kernels +from cunumpy import memory as memory +from cunumpy import mpi as mpi +from cunumpy import petsc as petsc +from cunumpy import profiling as profiling +from cunumpy import rng as rng +from cunumpy import xp as xp +from cunumpy._scipy_backend import scipy as scipy def to_numpy(array: Any) -> np.ndarray: ... def to_cupy(array: Any) -> Any: ... diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py index 5f55b5a..353c83a 100644 --- a/src/cunumpy/_algorithms.py +++ b/src/cunumpy/_algorithms.py @@ -8,7 +8,7 @@ import array_api_compat.numpy as np -from .xp import assert_same_backend, get_array_backend, get_array_module +from cunumpy.xp import assert_same_backend, get_array_backend, get_array_module def _integer_keys(keys: Any) -> tuple[Any, Any]: @@ -77,7 +77,7 @@ def cell_offsets(sorted_cells: Any, n_cells: int) -> Any: def _sum_kernel(dtype: Any) -> Any: name = np.dtype(dtype).name if name not in _SUM_KERNELS: - from ._cuda_kernel import CudaKernel + from cunumpy._cuda_kernel import CudaKernel ctype = "float" if name == "float32" else "double" _SUM_KERNELS[name] = CudaKernel( diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 344d88a..e022bed 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/src/cunumpy/_cuda_kernel.py @@ -1362,7 +1362,7 @@ def verify_layout( RuntimeError If CuPy is not available. """ - from .xp import cupy_available + from cunumpy.xp import cupy_available if not cupy_available(): raise RuntimeError("verify_layout() compiles a CUDA kernel and needs CuPy") @@ -1768,7 +1768,7 @@ def __host_args__(self) -> Any: "or set host_copies = True for a read-only host evaluation " "from host copies", ) - from .xp import to_numpy + from cunumpy.xp import to_numpy value = to_numpy(value) host_values.append(value) @@ -2093,7 +2093,7 @@ def debug_active(self) -> bool: """ if self._debug is not None: return self._debug - from ._device import get_cuda_debug + from cunumpy._device import get_cuda_debug return get_cuda_debug() @@ -2234,7 +2234,7 @@ def compile(self) -> Any: If CuPy or a GPU is not available. """ if self._raw_kernel is None: - from .xp import cupy_available + from cunumpy.xp import cupy_available if not cupy_available(): raise RuntimeError( @@ -2403,7 +2403,7 @@ def _opt_in_shared_memory(self, kernel: Any, shared_mem: int) -> None: attribute ``max_dynamic_shared_size_bytes``; it is set once (and again for a larger request) up to the device's opt-in limit. """ - from ._device import ( + from cunumpy._device import ( DEFAULT_SHARED_MEMORY_PER_BLOCK, max_shared_memory_per_block, ) diff --git a/src/cunumpy/_device.py b/src/cunumpy/_device.py index d8c1686..bb13907 100644 --- a/src/cunumpy/_device.py +++ b/src/cunumpy/_device.py @@ -9,8 +9,8 @@ import array_api_compat.numpy as np -from ._mpi_serial import local_rank -from .xp import array_backend, cupy_available +from cunumpy._mpi_serial import local_rank +from cunumpy.xp import array_backend, cupy_available def set_device(device_id: int) -> None: diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index ecbcb91..7a22014 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -34,16 +34,16 @@ import array_api_compat -from ._cuda_kernel import CudaKernel, _compile_in_threads -from ._kernel import ( +from cunumpy._cuda_kernel import CudaKernel, _compile_in_threads +from cunumpy._kernel import ( CompiledHostKernel, HostImplementations, PyccelKernel, resolve_host_args, ) -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _record -from .xp import get_backend +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _record +from cunumpy.xp import get_backend __all__ = ["Kernel", "KernelCatalog"] diff --git a/src/cunumpy/_emulation.py b/src/cunumpy/_emulation.py index 7ce3282..1cbf4d3 100644 --- a/src/cunumpy/_emulation.py +++ b/src/cunumpy/_emulation.py @@ -55,7 +55,7 @@ import numpy as np -from ._cuda_kernel import ( +from cunumpy._cuda_kernel import ( CudaKernel, CudaParameter, _scalar_checker, diff --git a/src/cunumpy/_fusion.py b/src/cunumpy/_fusion.py index 2299b4e..68541b8 100644 --- a/src/cunumpy/_fusion.py +++ b/src/cunumpy/_fusion.py @@ -32,7 +32,7 @@ def pressure(rho, T, gamma): import array_api_compat import numpy -from .xp import use_backend +from cunumpy.xp import use_backend __all__ = ["fuse"] diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index a05ca6e..1c74c1a 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -34,9 +34,9 @@ import array_api_compat import numpy as np -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _record -from .xp import _cupy_backend, _to_cupy, _to_numpy, is_gpu, to_cupy, to_numpy +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _record +from cunumpy.xp import _cupy_backend, _to_cupy, _to_numpy, is_gpu, to_cupy, to_numpy __all__ = ["CompiledHostKernel", "KernelArguments", "PyccelKernel", "resolve_host_args"] diff --git a/src/cunumpy/_mirror.py b/src/cunumpy/_mirror.py index d10294b..147735c 100644 --- a/src/cunumpy/_mirror.py +++ b/src/cunumpy/_mirror.py @@ -18,7 +18,7 @@ import numpy as np -from .xp import _cupy_backend, to_cupy +from cunumpy.xp import _cupy_backend, to_cupy __all__ = ["DeviceMirror"] diff --git a/src/cunumpy/_morton.py b/src/cunumpy/_morton.py index de39881..30e15ff 100644 --- a/src/cunumpy/_morton.py +++ b/src/cunumpy/_morton.py @@ -27,7 +27,7 @@ import numpy as np -from ._philox import _module +from cunumpy._philox import _module __all__ = [ "MAX_MORTON_LEVELS", diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index 6b1b158..9f9541f 100644 --- a/src/cunumpy/_mpi.py +++ b/src/cunumpy/_mpi.py @@ -11,12 +11,12 @@ import array_api_compat import array_api_compat.numpy as np -from ._mpi_serial import ( +from cunumpy._mpi_serial import ( _LOCAL_RANK_VARIABLES, # noqa: F401 - re-exported ) -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _describe, _record -from .xp import array_backend, cupy_available, to_numpy +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _describe, _record +from cunumpy.xp import array_backend, cupy_available, to_numpy _logger = logging.getLogger(__name__) @@ -43,7 +43,7 @@ def synchronize_for_mpi(*arrays: Any, stream: Any = None, event: Any = None) -> devices = {a.device.id for a in arrays if array_api_compat.is_cupy_array(a)} if not devices: return - from ._streams import HostEvent, HostStream + from cunumpy._streams import HostEvent, HostStream if isinstance(event, HostEvent) or isinstance(stream, HostStream): raise TypeError("device buffers require a CUDA producer stream or event") diff --git a/src/cunumpy/_profiling.py b/src/cunumpy/_profiling.py index e6402f7..586109a 100644 --- a/src/cunumpy/_profiling.py +++ b/src/cunumpy/_profiling.py @@ -9,7 +9,7 @@ from dataclasses import dataclass from types import ModuleType -from .xp import array_backend, synchronize +from cunumpy.xp import array_backend, synchronize def _nvtx_module() -> ModuleType | None: diff --git a/src/cunumpy/_random_streams.py b/src/cunumpy/_random_streams.py index 66daa3a..002ee86 100644 --- a/src/cunumpy/_random_streams.py +++ b/src/cunumpy/_random_streams.py @@ -34,7 +34,7 @@ import numpy as np -from .xp import get_backend +from cunumpy.xp import get_backend __all__ = ["BIT_GENERATORS", "RandomStreams", "random_streams"] diff --git a/src/cunumpy/_scipy_backend.py b/src/cunumpy/_scipy_backend.py index 1515563..64a0826 100644 --- a/src/cunumpy/_scipy_backend.py +++ b/src/cunumpy/_scipy_backend.py @@ -28,7 +28,7 @@ from types import ModuleType from typing import Any -from .xp import get_backend +from cunumpy.xp import get_backend __all__ = ["SUBMODULES", "ScipyNamespace", "scipy"] diff --git a/src/cunumpy/_staging.py b/src/cunumpy/_staging.py index 9faac25..b465c9b 100644 --- a/src/cunumpy/_staging.py +++ b/src/cunumpy/_staging.py @@ -34,8 +34,8 @@ import array_api_compat import numpy as np -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _record +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _record __all__ = ["HostStaging", "StagedCopy"] diff --git a/src/cunumpy/_streams.py b/src/cunumpy/_streams.py index 1d1f2db..706214f 100644 --- a/src/cunumpy/_streams.py +++ b/src/cunumpy/_streams.py @@ -4,7 +4,7 @@ from typing import Any -from .xp import get_backend +from cunumpy.xp import get_backend class HostEvent: diff --git a/src/cunumpy/algorithms.py b/src/cunumpy/algorithms.py index b974e03..c98421f 100644 --- a/src/cunumpy/algorithms.py +++ b/src/cunumpy/algorithms.py @@ -10,14 +10,14 @@ charge = xp.algorithms.segment_sum(q, cell, n_cells) """ -from ._algorithms import ( +from cunumpy._algorithms import ( SegmentPlan, cell_offsets, segment_boundaries, segment_sum, sort_by_key, ) -from ._morton import ( +from cunumpy._morton import ( MAX_MORTON_LEVELS, morton_decode, morton_encode, diff --git a/src/cunumpy/cuda/__init__.py b/src/cunumpy/cuda/__init__.py index 2aca624..412d548 100644 --- a/src/cunumpy/cuda/__init__.py +++ b/src/cunumpy/cuda/__init__.py @@ -24,7 +24,7 @@ use ``import cupy; cupy.cuda`` for CuPy's module. """ -from .._cuda_kernel import ( +from cunumpy._cuda_kernel import ( DEBUG_OPTIONS, CudaArguments, CudaKernel, @@ -41,7 +41,7 @@ resolve_includes, write_cuda_header, ) -from .._device import ( +from cunumpy._device import ( DEFAULT_SHARED_MEMORY_PER_BLOCK, bind_local_device, cuda_debug, @@ -56,7 +56,7 @@ set_device_for_rank, stream, ) -from .._streams import ( +from cunumpy._streams import ( HostEvent, HostStream, create_event, diff --git a/src/cunumpy/cuda_kernel.py b/src/cunumpy/cuda_kernel.py index cda7fac..516e70d 100644 --- a/src/cunumpy/cuda_kernel.py +++ b/src/cunumpy/cuda_kernel.py @@ -3,6 +3,6 @@ The public names are in :mod:`cunumpy.cuda`. """ -from ._deprecated import alias_module +from cunumpy._deprecated import alias_module alias_module(__name__, "cunumpy._cuda_kernel", "cunumpy.cuda") diff --git a/src/cunumpy/dispatch.py b/src/cunumpy/dispatch.py index ceb788d..8781b5b 100644 --- a/src/cunumpy/dispatch.py +++ b/src/cunumpy/dispatch.py @@ -3,6 +3,6 @@ The public names are in :mod:`cunumpy.kernels`. """ -from ._deprecated import alias_module +from cunumpy._deprecated import alias_module alias_module(__name__, "cunumpy._dispatch", "cunumpy.kernels") diff --git a/src/cunumpy/kernel.py b/src/cunumpy/kernel.py index 2780d36..d0b8717 100644 --- a/src/cunumpy/kernel.py +++ b/src/cunumpy/kernel.py @@ -3,6 +3,6 @@ The public names are in :mod:`cunumpy.kernels`. """ -from ._deprecated import alias_module +from cunumpy._deprecated import alias_module alias_module(__name__, "cunumpy._kernel", "cunumpy.kernels") diff --git a/src/cunumpy/kernel_testing.py b/src/cunumpy/kernel_testing.py index 4236582..50d3622 100644 --- a/src/cunumpy/kernel_testing.py +++ b/src/cunumpy/kernel_testing.py @@ -56,8 +56,8 @@ def test_parity(name, kernel): import array_api_compat import numpy as np -from . import _fake_cupy -from ._cuda_kernel import ( +from cunumpy import _fake_cupy +from cunumpy._cuda_kernel import ( CudaKernel, CudaParameter, CudaStructArguments, @@ -66,9 +66,9 @@ def test_parity(name, kernel): _split_top_level, _strip_comments, ) -from ._dispatch import Kernel -from ._emulation import emulate_cuda_kernel, emulation_compiler -from .xp import cupy_available, get_backend, to_numpy, use_backend +from cunumpy._dispatch import Kernel +from cunumpy._emulation import emulate_cuda_kernel, emulation_compiler +from cunumpy.xp import cupy_available, get_backend, to_numpy, use_backend # the pytest objects are created on first access, see __getattr__ __all__ = [ diff --git a/src/cunumpy/kernels.py b/src/cunumpy/kernels.py index 9149fb5..7f276c4 100644 --- a/src/cunumpy/kernels.py +++ b/src/cunumpy/kernels.py @@ -19,10 +19,10 @@ :mod:`cunumpy.kernel_testing`. """ -from ._cuda_kernel import PyccelStructArguments -from ._dispatch import Kernel, KernelCatalog -from ._fusion import fuse -from ._kernel import ( +from cunumpy._cuda_kernel import PyccelStructArguments +from cunumpy._dispatch import Kernel, KernelCatalog +from cunumpy._fusion import fuse +from cunumpy._kernel import ( HOST_IMPLEMENTATIONS, CompiledHostKernel, HostImplementations, diff --git a/src/cunumpy/memory.py b/src/cunumpy/memory.py index 27c91e3..7fb400b 100644 --- a/src/cunumpy/memory.py +++ b/src/cunumpy/memory.py @@ -16,7 +16,7 @@ :mod:`cunumpy.cuda`. """ -from ._mirror import DeviceMirror -from ._staging import HostStaging, StagedCopy +from cunumpy._mirror import DeviceMirror +from cunumpy._staging import HostStaging, StagedCopy __all__ = ["DeviceMirror", "HostStaging", "StagedCopy"] diff --git a/src/cunumpy/mpi.py b/src/cunumpy/mpi.py index b9d8206..f5a82b7 100644 --- a/src/cunumpy/mpi.py +++ b/src/cunumpy/mpi.py @@ -21,7 +21,7 @@ mpi4py is imported only by the functions that need it. """ -from ._mpi import ( +from cunumpy._mpi import ( MPIStaging, get_mpi_cuda_aware, mpi_buffer, @@ -30,7 +30,7 @@ set_mpi_cuda_aware, synchronize_for_mpi, ) -from ._mpi_serial import ( +from cunumpy._mpi_serial import ( OVERRIDE_VARIABLE, SerialComm, SerialMPI, diff --git a/src/cunumpy/profiling.py b/src/cunumpy/profiling.py index 7fb3ee7..662e797 100644 --- a/src/cunumpy/profiling.py +++ b/src/cunumpy/profiling.py @@ -13,8 +13,8 @@ step() """ -from ._profiling import Timing, nvtx_range, timed_region -from ._transfers import ( +from cunumpy._profiling import Timing, nvtx_range, timed_region +from cunumpy._transfers import ( TransferCounter, TransferEvent, assert_no_transfers, diff --git a/src/cunumpy/rng.py b/src/cunumpy/rng.py index 69ca9c1..4f31873 100644 --- a/src/cunumpy/rng.py +++ b/src/cunumpy/rng.py @@ -16,14 +16,19 @@ ``random`` module. """ -from ._philox import ( +from cunumpy._philox import ( philox4x32_10, philox_normal, philox_normal2, philox_uniform, philox_uniform2, ) -from ._random_streams import BIT_GENERATORS, RandomStreams, get_rng, random_streams +from cunumpy._random_streams import ( + BIT_GENERATORS, + RandomStreams, + get_rng, + random_streams, +) __all__ = [ "BIT_GENERATORS", diff --git a/src/cunumpy/testing.py b/src/cunumpy/testing.py index 8cd33d0..9f3593a 100644 --- a/src/cunumpy/testing.py +++ b/src/cunumpy/testing.py @@ -4,6 +4,6 @@ ``xp.testing.assert_allclose`` failed after any ``import cunumpy.testing``. """ -from ._deprecated import alias_module +from cunumpy._deprecated import alias_module alias_module(__name__, "cunumpy.kernel_testing", "cunumpy.kernel_testing") diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 5a56f1c..3da3621 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -13,12 +13,12 @@ import array_api_compat import array_api_compat.numpy as np -from ._transfers import _ACTIVE as _COUNTERS -from ._transfers import _describe, _record +from cunumpy._transfers import _ACTIVE as _COUNTERS +from cunumpy._transfers import _describe, _record if os.environ.get("CUNUMPY_FAKE_CUPY", "").strip().lower() in ("1", "true", "yes"): # tests without a GPU: a strict host stand-in for CuPy, see cunumpy._fake_cupy - from ._fake_cupy import install as _install_fake_cupy + from cunumpy._fake_cupy import install as _install_fake_cupy _install_fake_cupy()