nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/linalg/generic/_configuration/qualifiers.py

@@ -0,0 +1,494 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Matrix qualifier dataclasses for describing structured matrix types and their properties in
+linear algebra operations. Provides qualifiers for general, symmetric, hermitian,
+triangular, and diagonal matrices with associated metadata like fill modes, conjugate flags,
+and BLAS function abbreviations.
+"""
+
+import abc
+import dataclasses
+import typing
+
+import numpy as np
+import numpy.typing as npt
+
+import nvmath.bindings.cublas as cublas
+from nvmath._internal.templates import StatefulAPIOptions
+from nvmath.internal import utils
+
+FillMode: typing.TypeAlias = cublas.FillMode
+DiagType: typing.TypeAlias = cublas.DiagType
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class MatmulOptions(StatefulAPIOptions):
+    """A dataclass for providing options to a :class:`Matmul` object.
+
+    Attributes:
+        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).
+
+        blocking: A flag specifying the behavior of the stream-ordered functions and
+            methods. When ``blocking`` is `True`, the stream-ordered methods do not return
+            until the operation is complete. When ``blocking`` is ``"auto"``, the methods
+            return immediately when the inputs are on the GPU. The stream-ordered methods
+            always block when the operands are on the CPU to ensure that the user doesn't
+            inadvertently use the result before it becomes available. The default is
+            ``"auto"``.
+
+        inplace: Whether the matrix multiplication is performed in-place (operand C is
+            overwritten).
+
+        logger: Python Logger object. The root logger will be used if a
+            logger object is not provided.
+
+    .. seealso::
+       :class:`StatefulAPI`
+    """
+
+    inplace: bool = False
+
+
+MM_QUALIFIERS_DOCUMENTATION = {
+    #
+    "abbreviation": """\
+The two character abbreviation of the matrix qualifier.""".replace("\n", " "),
+    #
+    "conjugate": """\
+Whether the matrix is conjugate.""".replace("\n", " "),
+    #
+    "uplo": """\
+The :py:class:`~nvmath.bindings.cublas.FillMode` of the matrix. e.g. upper, lower...""".replace("\n", " "),
+    #
+    "diag": """\
+The :py:class:`~nvmath.bindings.cublas.DiagType` of the matrix. e.g. unit, non-unit...""".replace("\n", " "),
+    #
+    "incx": """\
+The direction to read the diagonal. +1 for forward; -1 for reverse.""".replace("\n", " "),
+    #
+}
+
+
+# We name this variable as foo_bar_dtype, so that the docstrings format correctly
+# Docstrings for custom dtypes are defined in docs/sphinx/conf.py
+matrix_qualifiers_dtype = np.dtype(
+    [
+        ("abbreviation", "U2"),
+        ("conjugate", np.bool_),
+        ("uplo", np.int_),
+        ("diag", np.int_),
+        ("incx", np.int_),
+    ]
+)
+
+"""A NumPy array of type :class:`matrix_qualifiers_dtype`.
+
+    .. seealso::
+        :class:`GeneralMatrixQualifier`, :class:`HermitianMatrixQualifier`,
+        :class:`SymmetricMatrixQualifier`, :class:`TriangularMatrixQualifier`
+        :class:`DiagonalMatrixQualifier`, :class:`matrix_qualifiers_dtype`
+"""
+MatrixQualifier: typing.TypeAlias = npt.NDArray
+
+
+class MatrixQualifierConstructor(abc.ABC):
+    abbreviation: typing.ClassVar[str]
+
+    def __init__(self):
+        msg = f"The {self.__class__.__name__} constructor should not be called. Use {self.__class__.__name__}.create() instead."
+        raise RuntimeError(msg)
+
+    @classmethod
+    def create(
+        cls,
+        conjugate: bool = False,
+        uplo: FillMode = FillMode.FULL,
+        diag: DiagType = DiagType.NON_UNIT,
+        incx: typing.Literal[-1, 1, 0] = 0,
+    ):
+        return np.array((cls.abbreviation, conjugate, uplo, diag, incx), dtype=matrix_qualifiers_dtype)
+
+    @classmethod
+    def is_valid(cls, other: MatrixQualifier) -> np.bool_:
+        """Return ``True`` if all elements of `other` are valid examples of the
+        :class:`matrix_qualifiers_dtype` constructed by this class."""
+        return np.all(other["abbreviation"] == cls.abbreviation) and np.bool_(
+            all(
+                n in other.dtype.names  # type: ignore[operator]
+                for n in ("conjugate",)
+            )
+        )
+
+    @classmethod
+    def to_string(cls, other: MatrixQualifier) -> str:
+        """Return a pretty string representation of `other`."""
+        return (
+            f"({other['abbreviation']}, conjugate={other['conjugate']}, "
+            f"uplo={FillMode(other['uplo']).name}, "
+            f"diag={DiagType(other['diag']).name}, incx={other['incx']})"
+        )
+
+
+@utils.docstring_decorator(MM_QUALIFIERS_DOCUMENTATION, skip_missing=False)
+class GeneralMatrixQualifier(MatrixQualifierConstructor):
+    """A class which constructs and validates :class:`matrix_qualifiers_dtype` for a general
+    rectangular matrix.
+
+    Examples:
+
+        >>> import numpy as np
+        >>> from nvmath.linalg import GeneralMatrixQualifier, matrix_qualifiers_dtype
+
+        Create a general matrix qualifier:
+
+        >>> GeneralMatrixQualifier.create()  # doctest: +ELLIPSIS
+        array(('ge', False, 2, 0, 0),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create a conjugate general matrix qualifier:
+
+        >>> GeneralMatrixQualifier.create(conjugate=True)  # doctest: +ELLIPSIS
+        array(('ge', True, 2, 0, 0),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create an array of general matrix qualifiers:
+
+        >>> np.full(
+        ...     2, GeneralMatrixQualifier.create(), dtype=matrix_qualifiers_dtype
+        ... )  # doctest: +ELLIPSIS
+        array([('ge', False, 2, 0, 0), ('ge', False, 2, 0, 0)],
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+    .. seealso::
+        :class:`GeneralMatrixQualifier`, :class:`HermitianMatrixQualifier`,
+        :class:`SymmetricMatrixQualifier`, :class:`TriangularMatrixQualifier`
+        :class:`DiagonalMatrixQualifier`, :class:`matrix_qualifiers_dtype`
+    """
+
+    abbreviation: typing.ClassVar[str] = "ge"
+
+    @classmethod
+    def create(  # type: ignore[override]
+        cls,
+        conjugate: bool = False,
+    ):
+        """Return a :class:`np.ndarray` of type :class:`matrix_qualifiers_dtype` whose
+        element describes a general matrix.
+
+        Args:
+            conjugate: {conjugate}
+
+        """
+        return super().create(conjugate=conjugate)
+
+
+@utils.docstring_decorator(MM_QUALIFIERS_DOCUMENTATION, skip_missing=False)
+class DiagonalMatrixQualifier(MatrixQualifierConstructor):
+    """A class which constructs and validates :class:`matrix_qualifiers_dtype` for a
+    diagonal matrix.
+
+    Examples:
+
+        >>> import numpy as np
+        >>> from nvmath.linalg import (
+        ...     DiagonalMatrixQualifier,
+        ...     GeneralMatrixQualifier,
+        ...     matrix_qualifiers_dtype,
+        ... )
+
+        Create a diagonal matrix qualifier:
+
+        >>> DiagonalMatrixQualifier.create()  # doctest: +ELLIPSIS
+        array(('dg', False, 2, 0, 1),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create a conjugate diagonal matrix qualifier:
+
+        >>> DiagonalMatrixQualifier.create(conjugate=True)  # doctest: +ELLIPSIS
+        array(('dg', True, 2, 0, 1),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create an array of matrix qualifiers with one general and one diagonal matrix:
+
+        >>> qualifiers = np.full(
+        ...     2,
+        ...     GeneralMatrixQualifier.create(),
+        ...     dtype=matrix_qualifiers_dtype,
+        ... )
+        >>> qualifiers[1] = DiagonalMatrixQualifier.create()
+        >>> qualifiers  # doctest: +ELLIPSIS
+        array([('ge', False, 2, 0, 0), ('dg', False, 2, 0, 1)],
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+    .. seealso::
+        :class:`GeneralMatrixQualifier`, :class:`HermitianMatrixQualifier`,
+        :class:`SymmetricMatrixQualifier`, :class:`TriangularMatrixQualifier`
+        :class:`DiagonalMatrixQualifier`, :class:`matrix_qualifiers_dtype`
+    """
+
+    abbreviation: typing.ClassVar[str] = "dg"
+
+    @classmethod
+    def create(  # type: ignore[override]
+        cls,
+        conjugate: bool = False,
+        incx: typing.Literal[-1, 1] = 1,
+    ):
+        """Return a :class:`np.ndarray` of type :class:`matrix_qualifiers_dtype` whose
+        element describes a diagonal matrix.
+
+        Args:
+            conjugate: {conjugate}
+
+            incx: {incx}
+        """
+        if incx not in (-1, 1):
+            raise ValueError(f"The 'incx' parameter must be '-1' or '1' not {incx}")
+        return super().create(conjugate, incx=incx)
+
+    @classmethod
+    def is_valid(cls, other):
+        return super().is_valid(other) and np.all(np.logical_or(other["incx"] == -1, other["incx"] == 1))
+
+
+@utils.docstring_decorator(MM_QUALIFIERS_DOCUMENTATION, skip_missing=False)
+class SquareMatrixQualifier(MatrixQualifierConstructor, abc.ABC):
+    @classmethod
+    def create(  # type: ignore[override]
+        cls,
+        conjugate: bool = False,
+        uplo: FillMode = FillMode.LOWER,
+        **kwargs,
+    ):
+        if uplo not in (FillMode.UPPER, FillMode.LOWER):
+            raise ValueError(f"The 'uplo' parameter must be 'UPPER' or 'LOWER', not {uplo}.")
+        return super().create(conjugate, uplo=uplo, **kwargs)
+
+    @classmethod
+    def is_valid(cls, other):
+        return super().is_valid(other) and np.all(
+            np.logical_or(other["uplo"] == FillMode.UPPER, other["uplo"] == FillMode.LOWER)
+        )
+
+
+@utils.docstring_decorator(MM_QUALIFIERS_DOCUMENTATION, skip_missing=False)
+class HermitianMatrixQualifier(SquareMatrixQualifier):
+    """A class which constructs and validates :class:`matrix_qualifiers_dtype` for a
+    hermitian matrix.
+
+    Examples:
+
+        >>> import numpy as np
+        >>> from nvmath.linalg import (
+        ...     HermitianMatrixQualifier,
+        ...     GeneralMatrixQualifier,
+        ...     matrix_qualifiers_dtype,
+        ... )
+
+        Create a hermitian matrix qualifier:
+
+        >>> HermitianMatrixQualifier.create()  # doctest: +ELLIPSIS
+        array(('he', False, 0, 0, 0),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create a conjugate hermitian matrix qualifier:
+
+        >>> HermitianMatrixQualifier.create(conjugate=True)  # doctest: +ELLIPSIS
+        array(('he', True, 0, 0, 0),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create an array of matrix qualifiers with one general and one hermitian matrix:
+
+        >>> qualifiers = np.full(
+        ...     2,
+        ...     GeneralMatrixQualifier.create(),
+        ...     dtype=matrix_qualifiers_dtype,
+        ... )
+        >>> qualifiers[1] = HermitianMatrixQualifier.create()
+        >>> qualifiers  # doctest: +ELLIPSIS
+        array([('ge', False, 2, 0, 0), ('he', False, 0, 0, 0)],
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+    .. seealso::
+        :class:`GeneralMatrixQualifier`, :class:`HermitianMatrixQualifier`,
+        :class:`SymmetricMatrixQualifier`, :class:`TriangularMatrixQualifier`
+        :class:`DiagonalMatrixQualifier`, :class:`matrix_qualifiers_dtype`
+    """
+
+    abbreviation: typing.ClassVar[str] = "he"
+
+    @classmethod
+    def create(  # type: ignore[override]
+        cls,
+        conjugate: bool = False,
+        uplo: FillMode = FillMode.LOWER,
+    ):
+        """Return a :class:`np.ndarray` of type :class:`matrix_qualifiers_dtype` whose
+        element describes a hermitian matrix.
+
+        Args:
+            conjugate: {conjugate}
+
+            uplo: {uplo}
+        """
+        return super().create(conjugate, uplo)
+
+
+@utils.docstring_decorator(MM_QUALIFIERS_DOCUMENTATION, skip_missing=False)
+class SymmetricMatrixQualifier(SquareMatrixQualifier):
+    """A class which constructs and validates :class:`matrix_qualifiers_dtype` for a
+    symmetric matrix.
+
+    Examples:
+
+        >>> import numpy as np
+        >>> from nvmath.linalg import (
+        ...     SymmetricMatrixQualifier,
+        ...     GeneralMatrixQualifier,
+        ...     matrix_qualifiers_dtype,
+        ... )
+
+        Create a symmetric matrix qualifier:
+
+        >>> SymmetricMatrixQualifier.create()  # doctest: +ELLIPSIS
+        array(('sy', False, 0, 0, 0),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create a conjugate symmetric matrix qualifier:
+
+        >>> SymmetricMatrixQualifier.create(conjugate=True)  # doctest: +ELLIPSIS
+        array(('sy', True, 0, 0, 0),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create an array of matrix qualifiers with one general and one symmetric matrix:
+
+        >>> qualifiers = np.full(
+        ...     2,
+        ...     GeneralMatrixQualifier.create(),
+        ...     dtype=matrix_qualifiers_dtype,
+        ... )
+        >>> qualifiers[1] = SymmetricMatrixQualifier.create()
+        >>> qualifiers  # doctest: +ELLIPSIS
+        array([('ge', False, 2, 0, 0), ('sy', False, 0, 0, 0)],
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+    .. seealso::
+        :class:`GeneralMatrixQualifier`, :class:`HermitianMatrixQualifier`,
+        :class:`SymmetricMatrixQualifier`, :class:`TriangularMatrixQualifier`
+        :class:`DiagonalMatrixQualifier`, :class:`matrix_qualifiers_dtype`
+    """
+
+    abbreviation: typing.ClassVar[str] = "sy"
+
+    @classmethod
+    def create(  # type: ignore[override]
+        cls,
+        conjugate: bool = False,
+        uplo: FillMode = FillMode.LOWER,
+    ):
+        """Return a :class:`np.ndarray` of type :class:`matrix_qualifiers_dtype` whose
+        element describes a symmetric matrix.
+
+        Args:
+            conjugate: {conjugate}
+
+            uplo: {uplo}
+        """
+        return super().create(conjugate, uplo)
+
+
+@utils.docstring_decorator(MM_QUALIFIERS_DOCUMENTATION, skip_missing=False)
+class TriangularMatrixQualifier(SquareMatrixQualifier):
+    """A class which constructs and validates :class:`matrix_qualifiers_dtype` for a
+    triangular matrix.
+
+    Examples:
+
+        >>> import numpy as np
+        >>> from nvmath.linalg import (
+        ...     TriangularMatrixQualifier,
+        ...     GeneralMatrixQualifier,
+        ...     matrix_qualifiers_dtype,
+        ... )
+
+        Create a triangular matrix qualifier:
+
+        >>> TriangularMatrixQualifier.create()  # doctest: +ELLIPSIS
+        array(('tr', False, 0, 0, 0),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create a conjugate triangular matrix qualifier:
+
+        >>> TriangularMatrixQualifier.create(conjugate=True)  # doctest: +ELLIPSIS
+        array(('tr', True, 0, 0, 0),
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+        Create an array of matrix qualifiers with one general and one triangular matrix:
+
+        >>> qualifiers = np.full(
+        ...     2,
+        ...     GeneralMatrixQualifier.create(),
+        ...     dtype=matrix_qualifiers_dtype,
+        ... )
+        >>> qualifiers[1] = TriangularMatrixQualifier.create()
+        >>> qualifiers  # doctest: +ELLIPSIS
+        array([('ge', False, 2, 0, 0), ('tr', False, 0, 0, 0)],
+              dtype=[('abbreviation', '<U2'), ('conjugate', '?'), ...
+
+    .. seealso::
+        :class:`GeneralMatrixQualifier`, :class:`HermitianMatrixQualifier`,
+        :class:`SymmetricMatrixQualifier`, :class:`TriangularMatrixQualifier`
+        :class:`DiagonalMatrixQualifier`, :class:`matrix_qualifiers_dtype`
+    """
+
+    abbreviation: typing.ClassVar[str] = "tr"
+
+    @classmethod
+    def create(  # type: ignore[override]
+        cls,
+        conjugate: bool = False,
+        uplo: FillMode = FillMode.LOWER,
+        diag: DiagType = DiagType.NON_UNIT,
+    ):
+        """Return a :class:`np.ndarray` of type :class:`matrix_qualifiers_dtype` whose
+        element describes a triangular matrix.
+
+        Args:
+            conjugate: {conjugate}
+
+            uplo: {uplo}
+
+            diag: {diag}
+        """
+        if diag not in (DiagType.UNIT, DiagType.NON_UNIT):
+            raise ValueError(f"The 'diag' parameter must be 'UNIT' or 'NON_UNIT', not {diag}.")
+        return super().create(conjugate, diag=diag, uplo=uplo)
+
+    @classmethod
+    def is_valid(cls, other):
+        return super().is_valid(other) and np.all(
+            np.logical_or(other["diag"] == DiagType.UNIT, other["diag"] == DiagType.NON_UNIT)
+        )
+
+
+def vector_to_square(
+    shape: typing.Sequence[int], strides: typing.Sequence[int], qualifier: MatrixQualifier
+) -> tuple[typing.Sequence[int], typing.Sequence[int]]:
+    """If `qualifier` is a DiagonalMatrixQualifier, convert `shape` and `stride` from a
+    vector to the equivalent square matrix."""
+    if DiagonalMatrixQualifier.is_valid(qualifier):
+        if len(shape) != 1:
+            msg = f"The shape of a diagonal matrix must be 1D; not {shape}."
+            raise ValueError(msg)
+        shape = tuple(shape) * 2
+        if len(strides) != 1:
+            msg = f"The strides of a diagonal matrix must be 1D; not {strides}."
+            raise ValueError(msg)
+        strides = (0, *strides)
+    return shape, strides