nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/tensor/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from ._configuration import *  # noqa: F403
+from .contract import *  # noqa: F403

nvmath/tensor/_configuration.py

@@ -0,0 +1,120 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "ContractionAlgo",
+    "ContractionAutotuneMode",
+    "ContractionJitMode",
+    "ContractionCacheMode",
+    "ContractionOptions",
+    "ExecutionCUDA",
+]
+
+
+from dataclasses import dataclass
+from logging import Logger
+from typing import Literal
+
+from nvmath._internal import templates
+from nvmath.bindings import cutensor
+from nvmath.memory import BaseCUDAMemoryManager
+
+ContractionAlgo = cutensor.Algo
+ContractionAutotuneMode = cutensor.AutotuneMode
+ContractionJitMode = cutensor.JitMode
+ContractionCacheMode = cutensor.CacheMode
+
+
+@dataclass
+class ContractionOptions:
+    """
+    A data class for providing options to the :class:`BinaryContraction` and
+    :class:`TernaryContraction` objects, or the wrapper functions
+    :func:`binary_contraction`and :func:`ternary_contraction`.
+
+    Attributes:
+        compute_type: The compute type to use for the contraction.
+            See :class:`~nvmath.tensor.ComputeDesc` for available compute types.
+        logger (logging.Logger): Python Logger object. The root logger will be used if a
+            logger object is not provided.
+
+        blocking: A flag specifying the behavior of the execution functions and methods,
+            such as :func:`binary_contraction` and :meth:`TernaryContraction.execute`.
+            When ``blocking`` is `True`, the execution methods do not return until the
+            operation is complete. When
+            ``blocking`` is ``"auto"``, the methods return immediately when the input tensor
+            is on the GPU. The execution methods always block when the input tensor is
+            on the CPU to ensure that the user doesn't inadvertently use the result
+            before it becomes available. The default is ``"auto"``.
+
+        handle: cuTensor library handle. A handle will be created if one is not provided.
+
+        allocator: An object that supports the :class:`BaseCUDAMemoryManager` protocol, used
+            to draw device memory. If an allocator is not provided, a memory allocator from
+            the library package will be used (:func:`torch.cuda.caching_allocator_alloc` for
+            PyTorch operands, :func:`cupy.cuda.alloc` otherwise).
+
+        memory_limit: Maximum memory available to the contraction operation.
+            It can be specified as a value (with optional suffix like K[iB], M[iB],
+            G[iB]) or as a percentage. The default is 80% of the device memory.
+
+        result_layout: The layout policy to use for the result: ``"auto"`` (default),
+            ``"C"`` (row-major), ``"F"`` (column-major), or ``"optimized"``.
+            With ``"auto"``, cuTENSOR execution requirements take precedence. When
+            those requirements permit, an optimized layout is chosen for this
+            contraction; otherwise, a compatible fallback layout is used. With
+            ``"optimized"``, strides are chosen directly from the contraction
+            expression and operand shapes to improve performance for this operation.
+            The chosen layout is local to this contraction and may not be optimal for
+            later operations that consume the result. This option is ignored if the
+            output operand is explicitly provided.
+
+    .. note::
+
+        - optimized ``result_layout`` currently only supports binary contraction via
+          :func:`binary_contraction` and :class:`BinaryContraction` when the addend
+          ``c`` is not specified.
+        - For both binary and ternary contractions, cuTENSOR may require the addend
+          and output operands to have the same strides. When an addend operand is
+          provided, explicitly specifying ``result_layout`` to ``"C"`` or ``"F"`` may
+          therefore fail at runtime if that layout is incompatible with the addend.
+          This constraint depends on the cuTENSOR library version.
+
+    .. seealso::
+        For supported compute types by data type, refer to the cuTENSOR documentation:
+
+        * `Binary contraction <https://docs.nvidia.com/cuda/cutensor/latest/api/cutensor.html#cutensorcreatecontraction>`_
+        * `Ternary contraction <https://docs.nvidia.com/cuda/cutensor/latest/api/cutensor.html#cutensorcreatecontractiontrinary>`_
+
+    """
+
+    compute_type: int | None = None
+    logger: Logger | None = None
+    blocking: Literal[True, "auto"] = "auto"
+    handle: int | None = None
+    allocator: BaseCUDAMemoryManager | None = None
+    memory_limit: int | str | None = r"80%"
+    result_layout: Literal["auto", "C", "F", "optimized"] = "auto"
+
+    def __post_init__(self):
+        if self.blocking not in (True, "auto"):
+            raise ValueError("The value specified for 'blocking' must be either True or 'auto'.")
+        if self.result_layout not in ("auto", "C", "F", "optimized"):
+            raise ValueError("The value specified for 'result_layout' must be either 'auto', 'C', 'F', or 'optimized'.")
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ExecutionCUDA(templates.ExecutionCUDA):
+    """
+    A data class for providing GPU execution options to the :class:`BinaryContraction` and
+    :class:`TernaryContraction` objects, or the wrapper functions
+    :func:`binary_contraction`and :func:`ternary_contraction`.
+
+    Attributes:
+        device_id: CUDA device ordinal (only used if the operand resides on the CPU). The
+            default value is 0.
+
+    """
+
+    pass

nvmath/tensor/_internal/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0

nvmath/tensor/_internal/cutensor_config_ifc.py

@@ -0,0 +1,279 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface classes to encapsulate low-level calls to get or set configuration information.
+"""
+
+__all__ = ["ContractionPlanPreference"]
+
+
+import numpy as np
+
+from nvmath.bindings import cutensor
+from nvmath.internal import utils
+
+
+class ContractionPlanPreference:
+    """
+    An interface to configure :meth:`nvmath.tensor.BinaryContraction.plan` and
+    :meth:`nvmath.tensor.TernaryContraction.plan`. The
+    current configuration can also be queried.
+    """
+
+    def __init__(self, contraction):
+        """
+        ctor for internal use only.
+        """
+        self._contraction = contraction
+        self._handle = self._contraction.handle
+
+        get_dtype = cutensor.get_plan_preference_attribute_dtype
+
+        self._autotune_mode = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.AUTOTUNE_MODE))
+        self._cache_mode = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.CACHE_MODE))
+        self._incremental_count = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.INCREMENTAL_COUNT))
+        self._algo = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.ALGO))
+        self._kernel_rank = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.KERNEL_RANK))
+        self._jit = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.JIT))
+        self._gpu_arch = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.GPU_ARCH))
+
+    def _check_valid_contraction_wrapper(self, *args, **kwargs):
+        if not self._contraction.valid_state:
+            raise RuntimeError("The ContractionPlanPreference object cannot be used after its contraction object is free'd.")
+
+    @staticmethod
+    def _get_scalar_attribute(contraction, name, attribute):
+        """
+        name      = cutensor PlanPreference enum for the attribute
+        attribute = numpy ndarray object into which the value is stored by cutensornet
+        """
+        assert contraction.plan_preference_ptr is not None, "Internal error"
+        cutensor.plan_preference_get_attribute(
+            contraction.handle, contraction.plan_preference_ptr, name, attribute.ctypes.data, attribute.dtype.itemsize
+        )
+
+    @staticmethod
+    def _set_scalar_attribute(contraction, name, attribute, value):
+        """
+        name      = cutensor PlanPreference enum for the attribute
+        attribute = numpy ndarray object into which the value is stored
+        value     = the value to set the the attribute to
+        """
+        assert contraction.plan_preference_ptr is not None, "Internal error"
+        attribute[0] = value
+        cutensor.plan_preference_set_attribute(
+            contraction.handle, contraction.plan_preference_ptr, name, attribute.ctypes.data, attribute.dtype.itemsize
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def autotune_mode(self):
+        """
+        Query the autotune mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorautotunemode-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.AUTOTUNE_MODE, self._autotune_mode
+        )
+        return self._autotune_mode.item()
+
+    @autotune_mode.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def autotune_mode(self, autotune_mode):
+        """
+        Set the autotune mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorautotunemode-t>`__
+        for more information.
+
+        Args:
+            autotune_mode: (nvmath.tensor.ContractionAutotuneMode) The autotune mode.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.AUTOTUNE_MODE, self._autotune_mode, autotune_mode
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def cache_mode(self):
+        """
+        Query the cache mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorcachemode-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.CACHE_MODE, self._cache_mode
+        )
+        return self._cache_mode.item()
+
+    @cache_mode.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def cache_mode(self, cache_mode):
+        """
+        Set the cache mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorcachemode-t>`__
+        for more information.
+
+        Args:
+            cache_mode: (nvmath.tensor.ContractionCacheMode) The cache mode.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.CACHE_MODE, self._cache_mode, cache_mode
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def incremental_count(self):
+        """
+        Query the incremental count. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorincrementalcount-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.INCREMENTAL_COUNT, self._incremental_count
+        )
+        return self._incremental_count.item()
+
+    @incremental_count.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def incremental_count(self, incremental_count):
+        """
+        Set the incremental count. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorincrementalcount-t>`__
+        for more information.
+
+        Args:
+            incremental_count: The incremental count.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.INCREMENTAL_COUNT, self._incremental_count, incremental_count
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def algo(self):
+        """
+        Query the algo. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensoralgo-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(self._contraction, cutensor.PlanPreferenceAttribute.ALGO, self._algo)
+        return self._algo.item()
+
+    @algo.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def algo(self, algo):
+        """
+        Set the algo. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensoralgo-t>`__
+        for more information.
+
+        Args:
+            algo: (nvmath.tensor.ContractionAlgo) The contraction algorithm.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.ALGO, self._algo, algo
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def kernel_rank(self):
+        """
+        Query the kernel rank. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorkernelrank-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.KERNEL_RANK, self._kernel_rank
+        )
+        return self._kernel_rank.item()
+
+    @kernel_rank.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def kernel_rank(self, kernel_rank):
+        """
+        Set the kernel rank. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorkernelrank-t>`__
+        for more information.
+
+        Args:
+            kernel_rank: The kernel rank.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.KERNEL_RANK, self._kernel_rank, kernel_rank
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def jit(self):
+        """
+        Query the jit compilation mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorjit-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(self._contraction, cutensor.PlanPreferenceAttribute.JIT, self._jit)
+        return self._jit.item()
+
+    @jit.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def jit(self, jit):
+        """
+        Set the jit compilation mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorjit-t>`__
+        for more information.
+
+        Args:
+            jit: (nvmath.tensor.ContractionJitMode) The JIT compilation mode.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(self._contraction, cutensor.PlanPreferenceAttribute.JIT, self._jit, jit)
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def gpu_arch(self):
+        """
+        Query the GPU architecture. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorplanpreferenceattribute-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.GPU_ARCH, self._gpu_arch
+        )
+        return self._gpu_arch.item()
+
+    @gpu_arch.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def gpu_arch(self, gpu_arch: int):
+        """
+        Set the GPU architecture. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorplanpreferenceattribute-t>`__
+        for more information.
+
+        Args:
+            gpu_arch: (int) The GPU architecture.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.GPU_ARCH, self._gpu_arch, gpu_arch
+        )

nvmath/tensor/_internal/cutensor_utils.py

@@ -0,0 +1,230 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+import weakref
+from collections.abc import Sequence
+
+from nvmath.bindings import cutensor
+from nvmath.internal.tensor_ifc import TensorHolder
+from nvmath.internal.typemaps import NAME_TO_DATA_TYPE
+
+__all__ = [
+    "compute_pointer_alignment",
+    "compute_strides",
+    "TensorDescriptor",
+]
+
+DEFAULT_ALIGNMENT_REQUIREMENT = 256
+
+
+def compute_pointer_alignment(ptr: int) -> int:
+    """
+    Compute the pointer alignment for the given pointer.
+
+    Args:
+        ptr: Pointer address as integer
+
+    Returns:
+        The alignment value (256, 128, 64, 32, 16, 8, 4, 2, or 1)
+    """
+    return 256 if ptr == 0 else min(ptr & -ptr, 256)
+
+
+def compute_strides(shape: Sequence[int], order: str = "C") -> Sequence[int]:
+    """
+    Compute the strides for the given shape and order.
+
+    Args:
+        shape: Shape of the tensor
+
+    Returns:
+        The strides
+    """
+    assert order in {"C", "F"}, f"Invalid order: {order}"
+    strides = [0] * len(shape)
+    stride = 1
+    if order == "C":
+        for axis in range(len(shape) - 1, -1, -1):
+            strides[axis] = stride
+            stride *= shape[axis]
+    elif order == "F":
+        for axis in range(len(shape)):
+            strides[axis] = stride
+            stride *= shape[axis]
+    else:
+        raise ValueError(f"Invalid order: {order}")
+    return strides
+
+
+def get_optimized_strides(
+    input_modes: Sequence[Sequence[str]],
+    output_modes: Sequence[str],
+    output_shape: Sequence[int],
+) -> Sequence[int]:
+    """
+    Get the optimized strides for a **binary** contraction output.
+
+    Output modes are split into three groups:
+        - hypermodes in both operands,
+        - un-contracted modes from operand A,
+        - un-contracted modes from operand B.
+
+    In the optimized/canonical output layout, the hyper modes are
+    placed at the beginning, followed by the group of modes with larger
+    combined size, then the other group of modes.
+
+    For example, "Pijab,Pabcd->Picjd" labels hypermode "P", A-only modes "i,j",
+    and B-only modes "c,d". The canonical mode order is either "Pijcd" or "Pcdij",
+    depending on which group has the larger combined size.
+
+    Args:
+        input_modes: Two sequences of mode labels (typically characters), one per operand.
+        output_modes: Output subscript aligned with ``output_shape``.
+        output_shape: Extent along each output mode.
+
+    Returns:
+        Strides in elements matching the memory order implied by ``output_modes``.
+    """
+    if len(input_modes) != 2:
+        raise ValueError("Only binary contraction is supported for optimized strides")
+    a_modes, b_modes = input_modes
+
+    size_dict = dict(zip(output_modes, output_shape, strict=True))
+
+    output_set = set(output_modes)
+    b_set = set(b_modes)
+
+    hyper_modes = [m for m in a_modes if m in output_set and m in b_set]
+    a_remaining_modes = [m for m in a_modes if m in output_set and m not in b_set]
+    b_remaining_modes = [m for m in b_modes if m in output_set and m not in hyper_modes]
+
+    combined_size_a = math.prod(size_dict[mode] for mode in a_remaining_modes)
+    combined_size_b = math.prod(size_dict[mode] for mode in b_remaining_modes)
+
+    if combined_size_a >= combined_size_b:
+        canonical_output_modes = hyper_modes + a_remaining_modes + b_remaining_modes
+    else:
+        canonical_output_modes = hyper_modes + b_remaining_modes + a_remaining_modes
+
+    canonical_output_shape = [size_dict[mode] for mode in canonical_output_modes]
+    canonical_strides = compute_strides(canonical_output_shape, order="C")
+
+    stride_map = dict(zip(canonical_output_modes, canonical_strides, strict=True))
+
+    # Permute strides back to the caller's output mode order.
+    optimized_strides = [stride_map[mode] for mode in output_modes]
+    return optimized_strides
+
+
+class TensorDescriptor:
+    """
+    A managed wrapper for cuTensor tensor descriptors with automatic resource cleanup.
+
+    This class wraps a cuTensor tensor descriptor (``cutensorTensorDescriptor_t``) and
+    automatically manages its lifetime. The underlying C descriptor is created on
+    initialization and automatically destroyed when this object is garbage collected.
+
+    Args:
+        handle: The cuTensor library handle.
+        extents: The extents (shape) of each mode of the tensor. Must have the same
+            length as ``strides``.
+        strides: The stride for each mode of the tensor, in elements of the base type.
+            Must have the same length as ``extents``.
+        dtype: The data type of the tensor as a string.
+        alignment: The alignment requirement in bytes for the base pointer that will
+            be used with this descriptor.
+
+    """
+
+    def __init__(
+        self,
+        handle: int,
+        extents: Sequence[int],
+        strides: Sequence[int],
+        dtype: str,
+        alignment: int,
+    ):
+        self._extents = extents
+        self._strides = strides
+        self._dtype = dtype
+        self._alignment = alignment
+        num_modes = len(extents)
+        assert num_modes == len(strides), "The number of modes and strides must be the same"
+        self._tensor_desc = cutensor.create_tensor_descriptor(
+            handle, num_modes, extents, strides, NAME_TO_DATA_TYPE[dtype], alignment
+        )
+        self._finalizer = weakref.finalize(self, cutensor.destroy_tensor_descriptor, self._tensor_desc)
+
+    @property
+    def ptr(self) -> int:
+        return self._tensor_desc
+
+    @property
+    def shape(self) -> Sequence[int]:
+        return self._extents
+
+    @property
+    def dtype(self) -> str:
+        return self._dtype
+
+    @property
+    def alignment(self) -> int:
+        return self._alignment
+
+    @property
+    def strides(self) -> Sequence[int]:
+        return self._strides
+
+    @classmethod
+    def from_tensor_holder(cls, handle: int, tensor: TensorHolder) -> "TensorDescriptor":
+        """
+        Create a TensorDescriptor from a TensorHolder.
+
+        Args:
+            handle: The cuTensor library handle.
+            tensor: The TensorHolder to create the descriptor from.
+
+        Returns:
+            A TensorDescriptor with properties derived from the TensorHolder.
+        """
+        return cls(
+            handle=handle,
+            extents=tensor.shape,
+            strides=tensor.strides,
+            dtype=tensor.dtype,
+            alignment=compute_pointer_alignment(tensor.data_ptr),
+        )
+
+    @classmethod
+    def from_metadata(
+        cls,
+        handle: int,
+        shape: Sequence[int],
+        dtype: str,
+        strides: Sequence[int] | None = None,
+        alignment: int = DEFAULT_ALIGNMENT_REQUIREMENT,
+    ) -> "TensorDescriptor":
+        """
+        Create a TensorDescriptor for the given shape, data type, strides, and alignment.
+
+        Args:
+            handle: The cuTensor library handle.
+            shape: The shape (extents) of the tensor.
+            dtype: The data type of the tensor as a string.
+            strides: The strides of the tensor. Defaults to C order strides.
+            alignment: The alignment requirement in bytes. Defaults to 256.
+
+        Returns:
+            A TensorDescriptor with the specified properties.
+        """
+        if strides is None:
+            strides = compute_strides(shape, order="C")
+        return cls(
+            handle=handle,
+            extents=shape,
+            strides=strides,
+            dtype=dtype,
+            alignment=alignment,
+        )

nvmath/tensor/_internal/data.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Defined in CPython:
+# https://github.com/python/cpython/blob/26bc2cc06128890ac89492eca20e83abe0789c1c/Objects/unicodetype_db.h#L6311-L6349
+
+__all__ = ["WHITESPACE_UNICODE"]
+
+_WHITESPACE_UNICODE_INTS = [
+    0x0009,
+    0x000A,
+    0x000B,
+    0x000C,
+    0x000D,
+    0x001C,
+    0x001D,
+    0x001E,
+    0x001F,
+    0x0020,
+    0x0085,
+    0x00A0,
+    0x1680,
+    0x2000,
+    0x2001,
+    0x2002,
+    0x2003,
+    0x2004,
+    0x2005,
+    0x2006,
+    0x2007,
+    0x2008,
+    0x2009,
+    0x200A,
+    0x2028,
+    0x2029,
+    0x202F,
+    0x205F,
+    0x3000,
+]
+
+
+WHITESPACE_UNICODE = "".join(chr(s) for s in _WHITESPACE_UNICODE_INTS)