nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/bindings/nvpl/fft.pxd

@@ -0,0 +1,100 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# This code was automatically generated with version 25.5, generator version 0.3.1.dev1471+g7122059e9. Do not modify it directly.
+
+from libc.stdint cimport intptr_t
+
+from .cyfft cimport *
+
+
+###############################################################################
+# Types
+###############################################################################
+
+
+
+
+cdef class Plan:
+    cdef intptr_t forward_handle
+    cdef intptr_t inverse_handle
+    cdef int precision
+    cdef int kind
+
+
+###############################################################################
+# Enum
+###############################################################################
+
+
+# The `UNSPECIFIED` value is not valid sign in NVPL's header,
+# but is used here to allow for planning C2C that can handle
+# both directions
+cpdef enum Sign:
+    UNSPECIFIED = 0
+    FORWARD = FFTW_FORWARD
+    INVERSE = FFTW_INVERSE
+
+
+cpdef enum PlannerFlags:
+    ESTIMATE = FFTW_ESTIMATE
+    MEASURE = FFTW_MEASURE
+    PATIENT = FFTW_PATIENT
+    EXHAUSTIVE = FFTW_EXHAUSTIVE
+    WISDOM_ONLY = FFTW_WISDOM_ONLY
+
+
+cpdef enum Precision:
+    FLOAT = 32
+    DOUBLE = 64
+
+
+cpdef enum Kind:
+    C2C = 1
+    C2R = 2
+    R2C = 3
+
+
+
+
+
+###############################################################################
+# Convenience wrappers/adapters
+###############################################################################
+
+cpdef int plan_with_nthreads(int precision, int num_threads) except -1
+cpdef int planner_nthreads(int precision) except? -1
+cpdef int init_threads(int precision) except -1
+cpdef int cleanup_threads(int precision) except -1
+cpdef Plan plan_many(int precision, int kind, int sign, int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags)
+cpdef int destroy(Plan plan) except -1
+cpdef int execute(Plan plan, intptr_t idata, intptr_t odata, int sign) except -1
+
+###############################################################################
+# Functions
+###############################################################################
+
+cpdef int get_version() except? -1
+cpdef intptr_t plan_many_c2c_double(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, int sign, unsigned flags) except? 0
+cpdef intptr_t plan_many_r2c_double(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags) except? 0
+cpdef intptr_t plan_many_c2r_double(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags) except? 0
+cpdef void execute_c2c_double(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef void execute_r2c_double(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef void execute_c2r_double(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef intptr_t plan_many_c2c_float(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, int sign, unsigned flags) except? 0
+cpdef intptr_t plan_many_r2c_float(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags) except? 0
+cpdef intptr_t plan_many_c2r_float(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags) except? 0
+cpdef void execute_c2c_float(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef void execute_r2c_float(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef void execute_c2r_float(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef int init_threads_double() except 0
+cpdef int init_threads_float() except 0
+cpdef void plan_with_nthreads_double(int nthreads) except*
+cpdef void plan_with_nthreads_float(int nthreads) except*
+cpdef int planner_nthreads_double() except? 0
+cpdef int planner_nthreads_float() except? 0
+cpdef void cleanup_threads_double() except*
+cpdef void cleanup_threads_float() except*
+cpdef void destroy_plan_double(intptr_t plan) except*
+cpdef void destroy_plan_float(intptr_t plan) except*

nvmath/bindings/nvpl/fft.pyi

@@ -0,0 +1,168 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import _cython_3_2_5
+import enum
+from typing import Any, Callable, ClassVar
+
+__pyx_capi__: dict
+__test__: dict
+cleanup_threads: _cython_3_2_5.cython_function_or_method
+cleanup_threads_double: _cython_3_2_5.cython_function_or_method
+cleanup_threads_float: _cython_3_2_5.cython_function_or_method
+destroy: _cython_3_2_5.cython_function_or_method
+destroy_plan_double: _cython_3_2_5.cython_function_or_method
+destroy_plan_float: _cython_3_2_5.cython_function_or_method
+execute: _cython_3_2_5.cython_function_or_method
+execute_c2c_double: _cython_3_2_5.cython_function_or_method
+execute_c2c_float: _cython_3_2_5.cython_function_or_method
+execute_c2r_double: _cython_3_2_5.cython_function_or_method
+execute_c2r_float: _cython_3_2_5.cython_function_or_method
+execute_r2c_double: _cython_3_2_5.cython_function_or_method
+execute_r2c_float: _cython_3_2_5.cython_function_or_method
+get_version: _cython_3_2_5.cython_function_or_method
+init_threads: _cython_3_2_5.cython_function_or_method
+init_threads_double: _cython_3_2_5.cython_function_or_method
+init_threads_float: _cython_3_2_5.cython_function_or_method
+plan_many: _cython_3_2_5.cython_function_or_method
+plan_many_c2c_double: _cython_3_2_5.cython_function_or_method
+plan_many_c2c_float: _cython_3_2_5.cython_function_or_method
+plan_many_c2r_double: _cython_3_2_5.cython_function_or_method
+plan_many_c2r_float: _cython_3_2_5.cython_function_or_method
+plan_many_r2c_double: _cython_3_2_5.cython_function_or_method
+plan_many_r2c_float: _cython_3_2_5.cython_function_or_method
+plan_with_nthreads: _cython_3_2_5.cython_function_or_method
+plan_with_nthreads_double: _cython_3_2_5.cython_function_or_method
+plan_with_nthreads_float: _cython_3_2_5.cython_function_or_method
+planner_nthreads: _cython_3_2_5.cython_function_or_method
+planner_nthreads_double: _cython_3_2_5.cython_function_or_method
+planner_nthreads_float: _cython_3_2_5.cython_function_or_method
+
+class FFTWError(Exception): ...
+
+class FFTWUnaligned(FFTWError): ...
+
+class Kind(enum.IntFlag):
+    __new__: ClassVar[Callable] = ...
+    C2C: ClassVar[Kind] = ...
+    C2R: ClassVar[Kind] = ...
+    R2C: ClassVar[Kind] = ...
+    _all_bits_: ClassVar[int] = ...
+    _boundary_: ClassVar[enum.FlagBoundary] = ...
+    _flag_mask_: ClassVar[int] = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _hashable_values_: ClassVar[list] = ...
+    _inverted_: ClassVar[None] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[int]] = ...
+    _singles_mask_: ClassVar[int] = ...
+    _unhashable_values_: ClassVar[list] = ...
+    _unhashable_values_map_: ClassVar[dict] = ...
+    _use_args_: ClassVar[bool] = ...
+    _value2member_map_: ClassVar[dict] = ...
+    __and__: ClassVar[Callable] = ...
+    __invert__: ClassVar[Callable] = ...
+    __or__: ClassVar[Callable] = ...
+    __rand__: ClassVar[Callable] = ...
+    __ror__: ClassVar[Callable] = ...
+    __rxor__: ClassVar[Callable] = ...
+    __xor__: ClassVar[Callable] = ...
+    def __format__(self, *args, **kwargs) -> str:
+        """Convert to a string according to format_spec."""
+
+class Plan:
+    @classmethod
+    def __init__(cls, *args, **kwargs) -> None:
+        """Create and return a new object.  See help(type) for accurate signature."""
+    def __reduce__(self) -> Any:
+        """Plan.__reduce__(self)"""
+
+class PlannerFlags(enum.IntFlag):
+    __new__: ClassVar[Callable] = ...
+    ESTIMATE: ClassVar[PlannerFlags] = ...
+    EXHAUSTIVE: ClassVar[PlannerFlags] = ...
+    MEASURE: ClassVar[PlannerFlags] = ...
+    PATIENT: ClassVar[PlannerFlags] = ...
+    WISDOM_ONLY: ClassVar[PlannerFlags] = ...
+    _all_bits_: ClassVar[int] = ...
+    _boundary_: ClassVar[enum.FlagBoundary] = ...
+    _flag_mask_: ClassVar[int] = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _hashable_values_: ClassVar[list] = ...
+    _inverted_: ClassVar[None] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[int]] = ...
+    _singles_mask_: ClassVar[int] = ...
+    _unhashable_values_: ClassVar[list] = ...
+    _unhashable_values_map_: ClassVar[dict] = ...
+    _use_args_: ClassVar[bool] = ...
+    _value2member_map_: ClassVar[dict] = ...
+    __and__: ClassVar[Callable] = ...
+    __invert__: ClassVar[Callable] = ...
+    __or__: ClassVar[Callable] = ...
+    __rand__: ClassVar[Callable] = ...
+    __ror__: ClassVar[Callable] = ...
+    __rxor__: ClassVar[Callable] = ...
+    __xor__: ClassVar[Callable] = ...
+    def __format__(self, *args, **kwargs) -> str:
+        """Convert to a string according to format_spec."""
+
+class Precision(enum.IntFlag):
+    __new__: ClassVar[Callable] = ...
+    DOUBLE: ClassVar[Precision] = ...
+    FLOAT: ClassVar[Precision] = ...
+    _all_bits_: ClassVar[int] = ...
+    _boundary_: ClassVar[enum.FlagBoundary] = ...
+    _flag_mask_: ClassVar[int] = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _hashable_values_: ClassVar[list] = ...
+    _inverted_: ClassVar[None] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[int]] = ...
+    _singles_mask_: ClassVar[int] = ...
+    _unhashable_values_: ClassVar[list] = ...
+    _unhashable_values_map_: ClassVar[dict] = ...
+    _use_args_: ClassVar[bool] = ...
+    _value2member_map_: ClassVar[dict] = ...
+    __and__: ClassVar[Callable] = ...
+    __invert__: ClassVar[Callable] = ...
+    __or__: ClassVar[Callable] = ...
+    __rand__: ClassVar[Callable] = ...
+    __ror__: ClassVar[Callable] = ...
+    __rxor__: ClassVar[Callable] = ...
+    __xor__: ClassVar[Callable] = ...
+    def __format__(self, *args, **kwargs) -> str:
+        """Convert to a string according to format_spec."""
+
+class Sign(enum.IntFlag):
+    __new__: ClassVar[Callable] = ...
+    FORWARD: ClassVar[Sign] = ...
+    INVERSE: ClassVar[Sign] = ...
+    UNSPECIFIED: ClassVar[Sign] = ...
+    _all_bits_: ClassVar[int] = ...
+    _boundary_: ClassVar[enum.FlagBoundary] = ...
+    _flag_mask_: ClassVar[int] = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _hashable_values_: ClassVar[list] = ...
+    _inverted_: ClassVar[None] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[int]] = ...
+    _singles_mask_: ClassVar[int] = ...
+    _unhashable_values_: ClassVar[list] = ...
+    _unhashable_values_map_: ClassVar[dict] = ...
+    _use_args_: ClassVar[bool] = ...
+    _value2member_map_: ClassVar[dict] = ...
+    __and__: ClassVar[Callable] = ...
+    __invert__: ClassVar[Callable] = ...
+    __or__: ClassVar[Callable] = ...
+    __rand__: ClassVar[Callable] = ...
+    __ror__: ClassVar[Callable] = ...
+    __rxor__: ClassVar[Callable] = ...
+    __xor__: ClassVar[Callable] = ...
+    def __format__(self, *args, **kwargs) -> str:
+        """Convert to a string according to format_spec."""

nvmath/bindings/nvshmem.pxd

@@ -0,0 +1,54 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# This code was automatically generated with version 3.1.7. Do not modify it directly.
+
+from libc.stdint cimport intptr_t
+
+from .cynvshmem cimport *
+
+
+###############################################################################
+# Types
+###############################################################################
+
+ctypedef nvshmem_team_config_v1 team_config
+ctypedef nvshmemx_uniqueid_args_v1 uniqueid_args
+ctypedef nvshmemx_init_args_v1 init_args
+ctypedef nvshmemx_init_attr_v1 init_attr
+
+ctypedef cudaStream_t Stream
+
+
+###############################################################################
+# Enum
+###############################################################################
+
+ctypedef _anon_enum0 __anon_enum0
+ctypedef _anon_enum1 __anon_enum1
+ctypedef _anon_enum2 __anon_enum2
+ctypedef nvshmemx_status _Status
+ctypedef flags _Flags
+
+
+###############################################################################
+# Functions
+###############################################################################
+
+cpdef int init_status() except? -1
+cpdef int my_pe() except? -1
+cpdef int n_pes() except? -1
+cpdef intptr_t malloc(size_t size) except? 0
+cpdef intptr_t calloc(size_t count, size_t size) except? 0
+cpdef intptr_t align(size_t alignment, size_t size) except? 0
+cpdef void free(intptr_t ptr) except*
+cpdef intptr_t ptr(intptr_t dest, int pe) except? 0
+cpdef void int_p(intptr_t dest, int value, int pe) except*
+cpdef int team_my_pe(int32_t team) except? -1
+cpdef void barrier_all_on_stream(intptr_t stream) except*
+cpdef void sync_all_on_stream(intptr_t stream) except*
+cpdef hostlib_init_attr(unsigned int flags, intptr_t attr)
+cpdef void hostlib_finalize() except*
+cpdef set_attr_uniqueid_args(int myrank, int nranks, intptr_t uniqueid, intptr_t attr)
+cpdef get_uniqueid(intptr_t uniqueid)

nvmath/bindings/nvshmem.pyi

@@ -0,0 +1,191 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import _cython_3_2_5
+import enum
+import numpy.dtypes
+from _typeshed import Incomplete
+from typing import Any, Callable, ClassVar
+
+STATUS_FULL_MPG: int
+STATUS_INVALID: int
+STATUS_IS_BOOTSTRAPPED: int
+STATUS_IS_INITIALIZED: int
+STATUS_LIMITED_MPG: int
+STATUS_NOT_INITIALIZED: int
+TEAMS_MIN: int
+TEAM_GPU_LEADERS: int
+TEAM_GPU_LEADERS_INDEX: int
+TEAM_INDEX_MAX: int
+TEAM_INVALID: int
+TEAM_NODE: int
+TEAM_NODE_INDEX: int
+TEAM_SAME_GPU: int
+TEAM_SAME_GPU_INDEX: int
+TEAM_SAME_MYPE_NODE: int
+TEAM_SAME_MYPE_NODE_INDEX: int
+TEAM_SHARED: int
+TEAM_SHARED_INDEX: int
+TEAM_WORLD: int
+TEAM_WORLD_INDEX: int
+__pyx_capi__: dict
+__test__: dict
+align: _cython_3_2_5.cython_function_or_method
+barrier_all_on_stream: _cython_3_2_5.cython_function_or_method
+calloc: _cython_3_2_5.cython_function_or_method
+check_status: _cython_3_2_5.cython_function_or_method
+free: _cython_3_2_5.cython_function_or_method
+get_uniqueid: _cython_3_2_5.cython_function_or_method
+hostlib_finalize: _cython_3_2_5.cython_function_or_method
+hostlib_init_attr: _cython_3_2_5.cython_function_or_method
+init_attr_dtype: numpy.dtypes.VoidDType
+init_status: _cython_3_2_5.cython_function_or_method
+int_p: _cython_3_2_5.cython_function_or_method
+malloc: _cython_3_2_5.cython_function_or_method
+my_pe: _cython_3_2_5.cython_function_or_method
+n_pes: _cython_3_2_5.cython_function_or_method
+ptr: _cython_3_2_5.cython_function_or_method
+set_attr_uniqueid_args: _cython_3_2_5.cython_function_or_method
+sync_all_on_stream: _cython_3_2_5.cython_function_or_method
+team_my_pe: _cython_3_2_5.cython_function_or_method
+uniqueid_dtype: numpy.dtypes.VoidDType
+
+class Flags(enum.IntEnum):
+    """See `flags`."""
+    __new__: ClassVar[Callable] = ...
+    INIT_MAX: ClassVar[Flags] = ...
+    INIT_THREAD_PES: ClassVar[Flags] = ...
+    INIT_WITH_MPI_COMM: ClassVar[Flags] = ...
+    INIT_WITH_SHMEM: ClassVar[Flags] = ...
+    INIT_WITH_UNIQUEID: ClassVar[Flags] = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _hashable_values_: ClassVar[list] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[int]] = ...
+    _unhashable_values_: ClassVar[list] = ...
+    _unhashable_values_map_: ClassVar[dict] = ...
+    _use_args_: ClassVar[bool] = ...
+    _value2member_map_: ClassVar[dict] = ...
+    def __format__(self, *args, **kwargs) -> str:
+        """Convert to a string according to format_spec."""
+
+class InitAttr:
+    """InitAttr()"""
+    mpi_comm: Incomplete
+    ptr: Incomplete
+    version: Incomplete
+    def __init__(self) -> Any:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+    def __reduce__(self):
+        """InitAttr.__reduce_cython__(self)"""
+    def __reduce_cython__(self) -> Any:
+        """InitAttr.__reduce_cython__(self)"""
+    def __setstate_cython__(self, __pyx_state) -> Any:
+        """InitAttr.__setstate_cython__(self, __pyx_state)"""
+
+class NvshmemError(Exception):
+    def __init__(self, status) -> Any:
+        """NvshmemError.__init__(self, status)"""
+    def __reduce__(self) -> Any:
+        """NvshmemError.__reduce__(self)"""
+
+class Status(enum.IntEnum):
+    """See `nvshmemx_status`."""
+    __new__: ClassVar[Callable] = ...
+    ERROR_COLLECTIVE_LAUNCH_FAILED: ClassVar[Status] = ...
+    ERROR_GPU_NOT_SELECTED: ClassVar[Status] = ...
+    ERROR_INTERNAL: ClassVar[Status] = ...
+    ERROR_INVALID_VALUE: ClassVar[Status] = ...
+    ERROR_NOT_SUPPORTED: ClassVar[Status] = ...
+    ERROR_OUT_OF_MEMORY: ClassVar[Status] = ...
+    ERROR_SENTINEL: ClassVar[Status] = ...
+    ERROR_SYMMETRY: ClassVar[Status] = ...
+    SUCCESS: ClassVar[Status] = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _hashable_values_: ClassVar[list] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[int]] = ...
+    _unhashable_values_: ClassVar[list] = ...
+    _unhashable_values_map_: ClassVar[dict] = ...
+    _use_args_: ClassVar[bool] = ...
+    _value2member_map_: ClassVar[dict] = ...
+    def __format__(self, *args, **kwargs) -> str:
+        """Convert to a string according to format_spec."""
+
+class UniqueId(uniqueid):
+    @classmethod
+    def __init__(cls, *args, **kwargs) -> None:
+        """Create and return a new object.  See help(type) for accurate signature."""
+    def __reduce__(self):
+        """UniqueId.__reduce_cython__(self)"""
+    def __reduce_cython__(self) -> Any:
+        """UniqueId.__reduce_cython__(self)"""
+    def __setstate_cython__(self, __pyx_state) -> Any:
+        """UniqueId.__setstate_cython__(self, __pyx_state)"""
+
+class uniqueid:
+    """uniqueid(size=1)
+
+    Empty-initialize an array of `nvshmemx_uniqueid_v1`.
+
+    The resulting object is of length `size` and of dtype `uniqueid_dtype`.
+    If default-constructed, the instance represents a single struct.
+
+    Args:
+        size (int): number of structs, default=1.
+
+
+    .. seealso:: `nvshmemx_uniqueid_v1`"""
+    ptr: Incomplete
+    version: Incomplete
+    def __init__(self, size=...) -> Any:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+    @staticmethod
+    def from_data(data) -> Any:
+        """uniqueid.from_data(data)
+
+        Create an uniqueid instance wrapping the given NumPy array.
+
+        Args:
+            data (_numpy.ndarray): a 1D array of dtype `uniqueid_dtype` holding the data."""
+    @staticmethod
+    def from_ptr(intptr_tptr, size_tsize=..., boolreadonly=...) -> Any:
+        """uniqueid.from_ptr(intptr_t ptr, size_t size=1, bool readonly=False)
+
+        Create an uniqueid instance wrapping the given pointer.
+
+        Args:
+            ptr (intptr_t): pointer address as Python :class:`int` to the data.
+            size (int): number of structs, default=1.
+            readonly (bool): whether the data is read-only (to the user). default is `False`."""
+    def __delitem__(self, other) -> None:
+        """Delete self[key]."""
+    def __eq__(self, other: object) -> bool:
+        """Return self==value."""
+    def __ge__(self, other: object) -> bool:
+        """Return self>=value."""
+    def __getitem__(self, index):
+        """Return self[key]."""
+    def __gt__(self, other: object) -> bool:
+        """Return self>value."""
+    def __int__(self) -> int:
+        """int(self)"""
+    def __le__(self, other: object) -> bool:
+        """Return self<=value."""
+    def __len__(self) -> int:
+        """Return len(self)."""
+    def __lt__(self, other: object) -> bool:
+        """Return self<value."""
+    def __ne__(self, other: object) -> bool:
+        """Return self!=value."""
+    def __reduce__(self):
+        """uniqueid.__reduce_cython__(self)"""
+    def __reduce_cython__(self) -> Any:
+        """uniqueid.__reduce_cython__(self)"""
+    def __setitem__(self, index, object) -> None:
+        """Set self[key] to value."""
+    def __setstate_cython__(self, __pyx_state) -> Any:
+        """uniqueid.__setstate_cython__(self, __pyx_state)"""

nvmath/device/__init__.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .common import *  # noqa: F403
+from .common_cuda import *  # noqa: F403
+from .cublasdx import *  # noqa: F403
+from .cublasdx_backend import *  # noqa: F403
+from .cufftdx import *  # noqa: F403
+from .cusolverdx import *  # noqa: F403
+from .types import *  # noqa: F403
+
+# isort: split
+
+# register models in compilers; must occur after imports above
+from .common import _HAS_NUMBA, _HAS_NUMBA_CUDA_MLIR
+
+if _HAS_NUMBA:
+    from . import (
+        cublasdx_numba,  # noqa: F401
+        cufftdx_numba,  # noqa: F401
+        cusolverdx_numba,  # noqa: F401
+    )
+    from .vector_types_numba import *  # noqa: F403
+
+if _HAS_NUMBA_CUDA_MLIR:
+    from numba_cuda_mlir.extending import refresh_registries
+
+    from . import (
+        common_numba_cuda_mlir,  # noqa: F401
+        cublasdx_numba_cuda_mlir,  # noqa: F401
+        cufftdx_numba_cuda_mlir,  # noqa: F401
+        cusolverdx_numba_cuda_mlir,  # noqa: F401
+    )
+
+    # Required for scenario when we compile some kernel
+    # and afterwards we import the device APIs
+    refresh_registries(include_uninitialized_cuda=False)

nvmath/device/_deprecated.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import sys
+from warnings import warn
+
+
+def _deprecated(
+    msg,
+    category: type[Warning] | None = DeprecationWarning,
+    stacklevel: int = 1,
+):
+    """Drop-in replacement for @warnings.deprecated(...)"""
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            warn(msg, category=category, stacklevel=stacklevel + 1)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+deprecated = _deprecated
+
+if (sys.version_info.major, sys.version_info.minor) >= (3, 13):
+    from warnings import deprecated as _warnings_deprecated  # type: ignore
+
+    deprecated = _warnings_deprecated