nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/linalg/generic/matmulmod.py

@@ -0,0 +1,2094 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "Matmul",
+    "matmul",
+]
+
+import dataclasses
+import logging
+import math
+from collections.abc import Generator, MutableSequence, Sequence
+from typing import Any, Final, Literal, TypeAlias
+
+import numpy as np
+
+from nvmath import memory
+from nvmath._internal import templates
+from nvmath._internal.utils import LoggerLike, get_addresses_of_elements
+from nvmath.bindings import cublas
+from nvmath.internal import formatters, tensor_wrapper, typemaps, utils
+from nvmath.internal._layout import StridedLayout
+from nvmath.internal.tensor_ifc_ndbuffer import NDBufferTensor
+from nvmath.linalg._internal.batch import BatchTraits
+from nvmath.linalg._internal.layout import BLASMatrixTraits, BLASMMTraitsView, InputMMTraits, check_extents, check_strides
+from nvmath.linalg.advanced.matmulmod import SHARED_MM_DOCUMENTATION
+from nvmath.linalg.generic._configuration import (
+    CACHED_LAYOUT_CHECKERS,
+    GeneralMatrixQualifier,
+    MatmulOptions,
+    MatrixQualifier,
+    matrix_qualifiers_dtype,
+    mm_layout_checker_getter,
+    select_blas_group_mm_function,
+    select_blas_mm_function,
+    vector_to_square,
+)
+from nvmath.linalg.generic._dtype import check_dtype
+
+AnyTensor: TypeAlias = utils.AnyTensor
+AnyStream: TypeAlias = utils.AnyStream
+SideMode: TypeAlias = cublas.SideMode
+FillMode: TypeAlias = cublas.FillMode
+DiagType: TypeAlias = cublas.DiagType
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class ExecutionCUDA(templates.ExecutionCUDA):
+    """
+    A data class for providing GPU execution options.
+
+    Attributes:
+        device_id: CUDA device ordinal (only used if the operand resides on the CPU). The
+            default value is 0.
+
+    .. seealso::
+        :class:`ExecutionCPU`, :class:`Matmul`, :func:`matmul`
+    """
+
+    pass
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class ExecutionCPU(templates.ExecutionCPU):
+    """
+    A data class for providing CPU execution options.
+
+    Attributes:
+        num_threads: The number of CPU threads used to execute the operation.
+                     If not specified, defaults to the number of CPU cores available to the
+                     process.
+
+    .. seealso::
+       :class:`ExecutionCUDA`, :class:`Matmul`, :func:`matmul`,
+       :func:`~nvmath.bindings.nvpl.blas.set_num_threads_local`
+    """
+
+    pass
+
+
+class InvalidMatmulState(Exception):
+    pass
+
+
+class _PrefixedLoggerAdapter(logging.LoggerAdapter):
+    def process(self, msg, kwargs):
+        return f"[{self.extra['prefix']}] {msg}", kwargs
+
+
+GENERIC_MM_DOCUMENTATION = SHARED_MM_DOCUMENTATION.copy()
+GENERIC_MM_DOCUMENTATION.update(
+    {
+        "a": SHARED_MM_DOCUMENTATION["a"].replace("tensor representing", "tensor or a sequence of tensors representing"),
+        #
+        "b": SHARED_MM_DOCUMENTATION["b"].replace("tensor representing", "tensor or a sequence of tensors representing"),
+        #
+        "c": SHARED_MM_DOCUMENTATION["c"].replace("tensor representing", "tensor or a sequence of tensors representing"),
+        #
+        "alpha": SHARED_MM_DOCUMENTATION["alpha"].replace("scale factor", "scale factor or a sequence of scale factors"),
+        #
+        "beta": SHARED_MM_DOCUMENTATION["beta"].replace("scale factor", "scale factor or a sequence of scale factors"),
+        #
+        "qualifiers": """\
+If desired, specify the matrix qualifiers as a :class:`numpy.ndarray` of
+:class:`~nvmath.linalg.generic.matrix_qualifiers_dtype` objects of length <= 3 corresponding to the operands `a`, `b`, and
+`c`. By default, :class:`GeneralMatrixQualifier` is assumed for each tensor. See
+:ref:`matrix-tensor-qualifiers` for the motivation behind qualifiers.""".replace("\n", " "),
+        #
+        "execution": """\
+Specify execution space options for the Matmul as a :class:`ExecutionCUDA` or :class:`ExecutionCPU` object. If not specified,
+the execution space will be selected to match operand's storage (in GPU or host memory), and the corresponding
+:class:`ExecutionCUDA` or :class:`ExecutionCPU` object will be default-constructed.""".replace("\n", " "),
+        #
+        "options": """\
+Specify options for the matrix multiplication as a :class:`MatmulOptions` object. Alternatively,
+a `dict` containing the parameters for the ``MatmulOptions`` constructor can also be provided.
+If not specified, the value will be set to the default-constructed ``MatmulOptions`` object.""".replace("\n", " "),
+        #
+        "result": """\
+The result of the specified matrix multiplication, which remains on the same device and belongs to the
+same package as the input operands.""".replace("\n", " "),
+        #
+        "semantics": """\
+        .. _semantics:
+
+        Inputs may be batched either **implicitly** and/or **explicitly**.
+
+        * An **implicitly-batched** input is a higher-dimensional :math:`N \\geq 3D` tensor.
+
+        * An **explicitly-batched** input is a Python sequence of tensors.
+
+        There may be only one explicitly-batched dimension, but there may be any number of implicitly-batched dimensions.
+        Explicitly-batched inputs are first broadcast to equal length along the explicitly-batched dimension,
+        then each explicit batch is treated as implicitly-batched inputs.
+
+        The semantics of the matrix multiplication follows :external:py:data:`numpy.matmul` semantics, with some restrictions.
+
+        * Broadcasting along implicitly-batched dimensions of `c` is not supported in this API, but may be supported in the
+          future. See the advanced API (:func:`nvmath.linalg.advanced.matmul`) for an API that supports broadcasting `c`.
+
+        In addition, the semantics for the fused matrix addition are described below:
+
+        * If arguments `a` and `b` are matrices, they are multiplied according to the rules of matrix multiplication.
+        * If argument `a` is 1-D, it is promoted to a matrix by prefixing ``1`` to its dimensions. After matrix
+          multiplication, the prefixed ``1`` is removed from the result's dimensions.
+        * If argument `b` is 1-D, it is promoted to a matrix by appending ``1`` to its dimensions. After matrix
+          multiplication, the appended ``1`` is removed from the result's dimensions.
+        * The operand for the matrix addition `c` must be the expected shape of the result of the matrix multiplication.
+
+""".strip(),
+        #
+        "reset_operands_unchecked": utils._reset_operand_unchecked_docstring(True, version_added="1.0"),
+    }
+)
+
+
+def _resolve_execution(
+    execution: ExecutionCPU | ExecutionCUDA | str | None,
+    operands_device_id: int | Literal["cpu"],
+) -> ExecutionCPU | ExecutionCUDA:
+    """Normalize the user-supplied ``execution`` argument to a concrete instance.
+
+    ``execution`` may be ``None``, the literal strings ``"cuda"``/``"cpu"``, or an
+    already-instantiated :class:`ExecutionCPU`/:class:`ExecutionCUDA`. The resolved
+    value's ``device_id`` is aligned with ``operands_device_id`` when both sides
+    pin a CUDA device.
+    """
+    match execution, operands_device_id:
+        case (None | "cuda", int()):
+            return ExecutionCUDA(device_id=operands_device_id)
+        case ("cuda", "cpu"):
+            return ExecutionCUDA()
+        case (None, "cpu") | ("cpu", _):
+            return ExecutionCPU()
+        case (ExecutionCUDA(), int()):
+            return dataclasses.replace(execution, device_id=operands_device_id)
+        case (ExecutionCPU(), _):
+            return execution
+        case (ExecutionCUDA(), "cpu"):
+            return execution
+        case _:
+            raise ValueError("Matmul.execution must be one of ExecutionCUDA, ExecutionCPU, None, 'cuda', or 'cpu'.")
+
+
+def _resolve_internal_operand_package(execution_name: str, operands_package: str) -> str:
+    """Map ``operands_package`` to the package name used internally for buffers/streams.
+
+    On CUDA execution, ``numpy`` operands are routed through ``cuda``; on CPU execution,
+    ``cupy`` operands are routed through ``cupy_host``. All other packages pass through.
+    """
+    if execution_name == "cuda":
+        return operands_package if operands_package != "numpy" else "cuda"
+    return operands_package if operands_package != "cupy" else "cupy_host"
+
+
+def _get_or_create_stream_maybe(
+    execution: ExecutionCPU | ExecutionCUDA,
+    operands_device_id: int | Literal["cpu"],
+    operands_package: str,
+    internal_op_package: str,
+    stream: utils.AnyStream,
+) -> tuple[utils.StreamHolder | None, utils.StreamHolder | None]:
+    """Return ``(exec_stream_holder, operand_stream_holder)`` for the given spaces.
+
+    The first holder is used for execution-space work (compute, workspace, output
+    buffers). The second is used for input-space data movement. When both are
+    non-``None`` and the execution and input spaces coincide, they are the same
+    holder.
+    """
+    if execution.name == "cuda":
+        h = utils.get_or_create_stream(execution.device_id, stream, internal_op_package)
+        return h, h
+    if isinstance(operands_device_id, int):
+        return None, utils.get_or_create_stream(operands_device_id, stream, operands_package)
+    return None, None
+
+
+def _realloc_if_needed_and_copy_to_mirror(
+    src_holder: utils.TensorHolder,
+    dest_holder: utils.TensorHolder,
+    dest_device_id: int | Literal["cpu"],
+    stream_holder: utils.StreamHolder | None,
+) -> None:
+    """Copy a source operand into its mirror on ``dest_device_id``.
+
+    If ``dest_holder.tensor`` is still alive, do an in-place copy into it. Otherwise
+    (typically after :meth:`release_operands`), allocate a new buffer via
+    ``src_holder.to()`` and bind it to ``dest_holder.tensor``.
+    """
+    if dest_holder.tensor is not None:
+        dest_holder.copy_(src_holder, stream_holder=stream_holder)
+        return
+    dest_holder.tensor = src_holder.to(dest_device_id, stream_holder).tensor
+
+
+class _ImplicitlyBatchedMatmul:
+    """A Matmul object that supports implicitly batched operands.
+
+    This class has the same public API as :class:`Matmul`, except that the arguments may
+    only be implicitly batched.
+
+    .. important::
+
+        Explicit batches are always defined as a Python Sequence of the same type and
+        implicit batches are always leading dimension of tensors. Otherwise, broadcasting
+        behavior would be ambiguous. For example, if some arguments are sequences of tensors
+        and other arguments are non-sequence tensors with batch dimensions.
+
+    .. seealso:: :func:`_zip_broadcast`, :class:`Matmul`, :class:`_ExplicitlyBatchedMatmul`
+
+    """
+
+    __slots__ = (
+        "_allocator",
+        "_batch_traits",
+        "_blocking",
+        "_call_prologue",
+        "_execution_device_id",
+        "_function",
+        "_input_traits",
+        "_internal_op_package",
+        "_logger",
+        "_memspace_eq_exespace",
+        "_mm_traits",
+        "_operands_backup",
+        "_operands_device_id",
+        "_operands_package",
+        "_operands",
+        "_qualifiers",
+        "_result_class",
+        "a_dtype_name",
+        "a_dtype",
+        "alpha",
+        "b_dtype_name",
+        "b_dtype",
+        "beta",
+        "c_dtype_name",
+        "c_dtype",
+        "execution",
+        "is_complex",
+        "lazy_conjugation",
+        "num_operands",
+        "operand_extents",
+        "operand_strides",
+        "options",
+        "scale_type_name",
+    )
+
+    options: Final[MatmulOptions]
+    _allocator: Final[memory.BaseCUDAMemoryManager | memory.BaseCUDAMemoryManagerAsync | None]
+    _blocking: Final[bool]
+    _logger: Final[LoggerLike]
+
+    # Metadata related to execution space
+    execution: Final[ExecutionCPU | ExecutionCUDA]
+    """A class which describes the execution space parameters."""
+    _internal_op_package: Final[str]
+    """The package of the operands in the execution space."""
+    _operands: MutableSequence[utils.TensorHolder]
+    """A copy of the operands in execution space."""
+    _result_class: Final[type[utils.TensorHolder]]
+    """The type of TensorHolder to use for the execution space result."""
+
+    # Metadata about the input/output tensors
+    _operands_backup: MutableSequence[utils.TensorHolder | None]
+    """A reference to original operands in their input space."""
+    _operands_device_id: Final[int | Literal["cpu"]]
+    """The device_id of the input space."""
+    _execution_device_id: Final[int | Literal["cpu"]]
+    """The device_id of the execution space (``"cpu"`` for CPU execution)."""
+    _memspace_eq_exespace: Final[bool]
+    """Cached ``_operands_device_id == _execution_device_id``."""
+    _operands_package: Final[str]
+    """The package of the operands in the input space."""
+
+    _call_prologue: Final[str]
+    """Stores a message for logging about blocking behavior"""
+
+    _input_traits: Final[InputMMTraits]
+    _batch_traits: Final[tuple[BatchTraits, BatchTraits, BatchTraits, BatchTraits]]
+    _qualifiers: Final[MatrixQualifier]
+
+    def __init__(
+        self,
+        a: AnyTensor,
+        b: AnyTensor,
+        /,
+        c: AnyTensor | None = None,
+        *,
+        alpha: float | complex | None = None,
+        beta: float | complex | None = None,
+        qualifiers: MatrixQualifier | None = None,
+        options: MatmulOptions,
+        execution: ExecutionCPU | ExecutionCUDA,
+        exec_stream_holder: utils.StreamHolder | None,
+        operand_stream_holder: utils.StreamHolder | None,
+        _allow_mixed_batch_order: bool = False,
+        _logger: LoggerLike | None = None,
+    ):
+        self.options = options
+        self._logger = _logger if _logger is not None else self.options.logger
+        self.execution = execution
+
+        assert not (c is None and options.inplace), "Internal Error: Matmul should have rejected inplace=True with c=None."
+
+        a = tensor_wrapper.wrap_operand(a)
+        check_dtype(a.dtype, "A")
+        check_extents(a.shape, "A")
+        check_strides(a.strides, "A")
+
+        b = tensor_wrapper.wrap_operand(b)
+        check_dtype(b.dtype, "B")
+        check_extents(b.shape, "B")
+        check_strides(b.strides, "B")
+        operands = [a, b]
+
+        self.num_operands = 2
+        if c is not None:
+            self.num_operands = 3
+            c = tensor_wrapper.wrap_operand(c)
+            check_dtype(c.dtype, "C")
+            check_extents(c.shape, "C")
+            check_strides(c.strides, "C")
+            operands.append(c)
+
+        self._operands_device_id = utils.get_operands_device_id(operands)
+        self._execution_device_id = getattr(self.execution, "device_id", "cpu")
+        self._memspace_eq_exespace = self._operands_device_id == self._execution_device_id
+        self._operands_package = utils.get_operands_package(operands)
+        self._internal_op_package = _resolve_internal_operand_package(self.execution.name, self._operands_package)
+
+        operands_backup: list[utils.TensorHolder | None] = [None] * len(operands)
+        for i in range(len(operands)):
+            # Copy the operand to execution_space's device if needed.
+            operands[i], operands_backup[i] = templates.copy_operand_perhaps(
+                None,
+                operands[i],
+                operand_stream_holder,
+                self._execution_device_id,
+                self._operands_device_id,
+            )
+        self._operands = operands
+        self._operands_backup = operands_backup
+
+        # The result's package and device.
+        self._result_class = self._operands[0].__class__
+
+        # Set blocking or non-blocking behavior.
+        self._blocking = self.options.blocking != "auto" or self._operands_device_id == "cpu" or self.execution.name == "cpu"
+        if self._blocking:
+            call_prologue = "This call is blocking and will return only after the operation is complete."
+        else:
+            call_prologue = (
+                "This call is non-blocking and will return immediately after the operation is launched on the device."
+            )
+        self._call_prologue = call_prologue
+
+        # Set memory allocator.
+        allocator: memory.BaseCUDAMemoryManager | memory.BaseCUDAMemoryManagerAsync | None
+        match self.execution:
+            case ExecutionCUDA():
+                allocator = (
+                    memory._MEMORY_MANAGER[self._internal_op_package](self.execution.device_id, self._logger)
+                    if self.options.allocator is None
+                    else self.options.allocator
+                )
+            case ExecutionCPU() | _:
+                allocator = None  # currently, the nvpl/fftw does not support custom workspace allocation
+        self._allocator = allocator
+
+        self._logger.info(f"The data type of operand A is '{a.dtype}', and that of operand B is '{b.dtype}'.")
+        if c is not None:
+            self._logger.info(f"The data type of operand C is '{c.dtype}'.")
+
+        if self.options.inplace:
+            self._logger.info("The operation will be performed inplace with operand C.")
+            assert c is not None  # guaranteed by the inplace-with-c assertion above
+
+            # Inplace writes back into C, so the C layout must be unique.
+            # Otherwise distinct logical indices alias the same memory cell,
+            # producing overlapping writes / undefined behavior.
+            if not StridedLayout(c.shape, c.strides, itemsize=1).is_unique:
+                raise ValueError(
+                    f"The layout for C with shape = {c.shape} and strides = {c.strides} is not "
+                    f"supported for inplace operations: the index-to-memory mapping is not "
+                    f"injective, which would cause overlapping writes into the same memory."
+                )
+
+        if c is not None and beta is None:
+            raise ValueError("A value for beta must be provided if operand C is provided.")
+
+        assert self.num_operands == 2 or self.num_operands == 3, "Internal Error."
+
+        if a.dtype != b.dtype or (c is not None and a.dtype != c.dtype):
+            raise ValueError(
+                "Unsupported combination of dtypes. "
+                f"A ({a.dtype}), B ({b.dtype}), and C ({getattr(c, 'dtype', None)}) must all have the same dtype."
+            )
+        # Determine the data types for a and b.
+        self.a_dtype = typemaps.NAME_TO_DATA_TYPE[a.dtype]
+        self.b_dtype = typemaps.NAME_TO_DATA_TYPE[b.dtype]
+        self.a_dtype_name = a.dtype
+        self.b_dtype_name = b.dtype
+
+        self.is_complex = "complex" in self.a_dtype_name or "complex" in self.b_dtype_name
+
+        # Determine the data types for c.
+        if c is None:
+            self.c_dtype = self.a_dtype
+        else:
+            self.c_dtype = typemaps.NAME_TO_DATA_TYPE[c.dtype]
+        self.c_dtype_name = typemaps.DATA_TYPE_TO_NAME[self.c_dtype]
+
+        self._logger.info(f"The data type for the result D is '{self.c_dtype_name}'.")
+
+        self.scale_type_name = self.a_dtype_name
+
+        # Set alpha and beta.
+        self.alpha = np.zeros((1,), dtype=self.scale_type_name)
+        try:
+            self.alpha[0] = alpha if alpha is not None else 1
+        except (ValueError, TypeError) as e:
+            raise ValueError(f"The value provided for alpha {alpha} is not convertible to dtype '{self.alpha.dtype}'.") from e
+
+        self.beta = np.zeros((1,), dtype=self.scale_type_name)
+        if beta is not None and c is None:
+            self._logger.warning(f"Matmul: The provided beta value {beta} is ignored since operand C is not specified.")
+        try:
+            self.beta[0] = beta if beta is not None and c is not None else 0
+        except (ValueError, TypeError) as e:
+            raise ValueError(f"The value provided for beta {beta} is not convertible to dtype '{self.beta.dtype}'.") from e
+
+        if qualifiers is None:
+            new_qualifiers = np.empty(3, dtype=matrix_qualifiers_dtype)
+            new_qualifiers[:] = GeneralMatrixQualifier.create()
+        else:
+            if not ((len(qualifiers) == 3) or (len(qualifiers) == 2 and c is None)):
+                raise ValueError("The number of MatrixQualifiers must match the number of operands.")
+            new_qualifiers = np.empty(3, dtype=matrix_qualifiers_dtype)
+            new_qualifiers[:2] = qualifiers[:2]
+            new_qualifiers[2] = GeneralMatrixQualifier.create() if len(qualifiers) < 3 else qualifiers[2]
+        self._qualifiers = new_qualifiers
+        self._logger.info(
+            f"The matrix multiplication qualifiers are:\n"
+            f"    A = {GeneralMatrixQualifier.to_string(self._qualifiers[0])}\n"
+            f"    B = {GeneralMatrixQualifier.to_string(self._qualifiers[1])}\n"
+            f"    C = {GeneralMatrixQualifier.to_string(self._qualifiers[2])}"
+        )
+
+        # Set qualifiers based on torch lazy conjugation flag if not provided.
+        self._qualifiers[0]["conjugate"] = self._qualifiers[0]["conjugate"] ^ self._operands[0].is_conjugate
+        self._qualifiers[1]["conjugate"] = self._qualifiers[1]["conjugate"] ^ self._operands[1].is_conjugate
+        self.lazy_conjugation = (self._operands[0].is_conjugate, self._operands[1].is_conjugate, False)
+        if c is not None:
+            self._qualifiers[2]["conjugate"] = self._qualifiers[2]["conjugate"] ^ self._operands[2].is_conjugate
+        if self._qualifiers[2]["conjugate"]:
+            raise ValueError("The conjugate flag is currently not supported for operand C.")
+
+        # Capture operand extents and strides for consistency check when resetting operands.
+        self.operand_extents = tuple(o.shape for o in self._operands)
+        self.operand_strides = tuple(o.strides for o in self._operands)
+
+        # Create operand layouts.
+        a_layout = BLASMatrixTraits(
+            self.a_dtype,
+            *vector_to_square(
+                self._operands[0].shape[-2:],
+                self._operands[0].strides[-2:],
+                self._qualifiers[0],
+            ),
+            is_conjugate=bool(self._qualifiers[0]["conjugate"]),
+            is_transpose=False,
+            is_lower=self._qualifiers[0]["uplo"] == FillMode.LOWER,
+        )
+        self._logger.info("Operand A has traits of %s", a_layout)
+        b_layout = BLASMatrixTraits(
+            self.b_dtype,
+            *vector_to_square(
+                self._operands[1].shape[-2:],
+                self._operands[1].strides[-2:],
+                self._qualifiers[1],
+            ),
+            is_conjugate=bool(self._qualifiers[1]["conjugate"]),
+            is_transpose=False,
+            is_lower=self._qualifiers[1]["uplo"] == FillMode.LOWER,
+        )
+        self._logger.info("Operand B has traits of %s", b_layout)
+        c_layout = (
+            None
+            if c is None
+            else BLASMatrixTraits(
+                self.c_dtype,
+                *vector_to_square(
+                    self._operands[2].shape[-2:],
+                    self._operands[2].strides[-2:],
+                    self._qualifiers[2],
+                ),
+                is_conjugate=bool(self._qualifiers[2]["conjugate"]),
+                is_transpose=False,
+                is_lower=self._qualifiers[2]["uplo"] == FillMode.LOWER,
+            )
+        )
+        self._logger.info("Operand C has traits of %s", c_layout)
+
+        # Get the operation traits.
+        self._input_traits = InputMMTraits.from_layouts(a_layout, b_layout, c_layout, self.options.inplace, self._logger)
+
+        a_batch = BatchTraits.from_full_shape_and_strides(
+            self._operands[0].shape,
+            self._operands[0].strides,
+            num_trailing_dims=2,
+            overlap_allowed=True,
+        )
+        b_batch = BatchTraits.from_full_shape_and_strides(
+            self._operands[1].shape,
+            self._operands[1].strides,
+            num_trailing_dims=2,
+            overlap_allowed=True,
+        )
+        c_batch = (
+            BatchTraits(shape=(), strides=())
+            if c is None
+            else BatchTraits.from_full_shape_and_strides(
+                self._operands[2].shape,
+                self._operands[2].strides,
+                num_trailing_dims=2,
+                overlap_allowed=False,
+            )
+        )
+        # BatchTraits overloads the * operator with a batch combination behavior
+        abc_batch, axis_order = a_batch.combine(
+            b_batch.combine(c_batch, allow_mixed_axis_order=_allow_mixed_batch_order),
+            allow_mixed_axis_order=_allow_mixed_batch_order,
+        )
+        d_batch = (
+            c_batch
+            if self.options.inplace
+            else BatchTraits.from_batch_shape_and_size(
+                shape=abc_batch,
+                axis_order=axis_order,
+                batch_stride=math.prod(self._input_traits.d_layout_traits.shape),
+            )
+        )
+        self._logger.info("Operand A has batch traits of %s.", a_batch)
+        self._logger.info("Operand B has batch traits of %s.", b_batch)
+        self._logger.info("Operand C has batch traits of %s.", c_batch)
+        self._logger.info("Result D has batch traits of %s.", d_batch)
+        self._batch_traits = (a_batch, b_batch, c_batch, d_batch)
+
+    def _get_or_create_stream_maybe(
+        self, stream: utils.AnyStream
+    ) -> tuple[utils.StreamHolder | None, utils.StreamHolder | None]:
+        return _get_or_create_stream_maybe(
+            self.execution,
+            self._operands_device_id,
+            self._operands_package,
+            self._internal_op_package,
+            stream,
+        )
+
+    def plan(self, *, stream: AnyStream | int | None = None, log_info=False) -> None:
+        # NOTE: timing context called from Matmul.plan to avoid redundancy
+        mm_layout_checker = mm_layout_checker_getter(self._qualifiers)
+
+        mm_traits = BLASMMTraitsView.from_input_traits(
+            self._input_traits,
+            mm_layout_checker,
+            self._logger,
+            lookup_table_table=CACHED_LAYOUT_CHECKERS,
+        )
+
+        self._mm_traits = mm_traits
+
+        self._function, function_name = select_blas_mm_function(
+            (*self._batch_traits[:2], self._batch_traits[3]),
+            mm_traits,
+            self._qualifiers,
+            self._logger,
+            self.execution,
+        )
+
+        if log_info:
+            flop_count = 2 * mm_traits.M * mm_traits.N * mm_traits.K
+
+            self._logger.info(f"The plan found 1 suitable algorithm named {function_name}.")
+            self._logger.info(f"The base matrix multiplication FLOP count is {formatters.FLOPSStr(flop_count, 'FLOP')}.")
+
+    def _check_and_set_operand(
+        self,
+        new_operand: utils.TensorHolder,
+        operand_name: str,
+        stream_holder: utils.StreamHolder | None,
+        *,
+        operand_index: int,
+        dtype: str | None = None,
+        extents: Sequence[int] | None = None,
+        strides: Sequence[int] | None = None,
+    ):
+        """
+        Check to make sure that the provided operand is consistent with the one it's
+        updating, and update it.
+        """
+        # Make sure that the data type and extents match.
+        utils.check_attribute_match(dtype, new_operand.dtype, "data type")
+        utils.check_attribute_match(extents, new_operand.shape, "extents")
+
+        # Package must be the same to preserve stream ordering
+        if self._operands_package != new_operand.name:
+            raise TypeError(
+                f"Library package mismatch: The operand {operand_name} must from the same package ({new_operand.name}) "
+                f"as the original operand ({self._operands_package})."
+            )
+
+        if self._operands_device_id != new_operand.device_id:
+            raise ValueError(
+                f"The operand {operand_name} must be on the same device ({new_operand.device_id}) as the original operand "
+                f"({self._operands_device_id})."
+            )
+
+        # Conjugate flag of the provided operands must match the original qualifiers
+        if self.lazy_conjugation[operand_index] != new_operand.is_conjugate:
+            raise ValueError(f"The provided operand {operand_name} has different conjugate flag than the original operand")
+
+        # Wrappers in self._operands / self._operands_backup are preserved for the
+        # lifetime of this object (release_operands clears only their .tensor); rebind
+        # in same-space and use the realloc-or-copy helper in cross-space, mirroring
+        # the unchecked path so both reset flavors maintain that invariant.
+        if self._memspace_eq_exespace:
+            self._operands[operand_index].tensor = new_operand.tensor
+        else:
+            _realloc_if_needed_and_copy_to_mirror(
+                new_operand,
+                self._operands[operand_index],
+                self._execution_device_id,
+                stream_holder,
+            )
+            self._operands_backup[operand_index] = new_operand
+
+        # After the rebind/copy above, check that the strides of the buffer the matmul
+        # will actually read (the user's tensor in same-space, the mirror in cross-space)
+        # still match what was planned for at construction.
+        # FIXME: a stride mismatch here leaves _operands[operand_index] already mutated;
+        # if the caller catches this they can observe a half-applied reset.
+        utils.check_attribute_match(strides, self._operands[operand_index].strides, "strides")
+
+        self._logger.info(f"Operand '{operand_name}' has been reset to the new value.")
+
+        return
+
+    def reset_operands(
+        self,
+        *,
+        a=None,
+        b=None,
+        c=None,
+        alpha=None,
+        beta=None,
+        stream: AnyStream | int | None = None,
+        _operands_released=True,
+    ):
+        if c is not None and self.num_operands == 2:
+            raise ValueError(
+                "The matrix multiplication problem specification does not include operand C, so it cannot be reset."
+            )
+
+        # If operands have been released, all required operands must be provided
+        if _operands_released:
+            all_provided = a is not None and b is not None and (self.num_operands != 3 or c is not None)
+            if not all_provided:
+                raise ValueError(
+                    "After release_operands(), all required operands must be provided to reset_operands(). "
+                    f"Required: a, b{', c' if self.num_operands == 3 else ''}"
+                )
+
+        # Update alpha.
+        if alpha is not None:
+            try:
+                self.alpha[0] = alpha
+            except (ValueError, TypeError) as e:
+                raise ValueError(
+                    f"The value provided for alpha {alpha} is not convertible to dtype '{self.alpha.dtype}'."
+                ) from e
+
+        # Update beta.
+        if beta is not None:
+            if self.num_operands == 2:
+                self._logger.warning("Matmul: The provided beta value %s is ignored since operand C is not specified.", beta)
+            else:
+                try:
+                    self.beta[0] = beta
+                except (ValueError, TypeError) as e:
+                    raise ValueError(
+                        f"The value provided for beta {beta} is not convertible to dtype '{self.beta.dtype}'."
+                    ) from e
+
+        _, operand_stream_holder = self._get_or_create_stream_maybe(stream)
+
+        # Reset the provided operands.
+        if a is not None:
+            a = tensor_wrapper.wrap_operand(a)
+            index = 0
+            self._check_and_set_operand(
+                a,
+                "A",
+                operand_stream_holder,
+                operand_index=index,
+                dtype=self.a_dtype_name,
+                extents=self.operand_extents[index],
+                strides=self.operand_strides[index],
+            )
+
+        if b is not None:
+            b = tensor_wrapper.wrap_operand(b)
+            index = 1
+            self._check_and_set_operand(
+                b,
+                "B",
+                operand_stream_holder,
+                operand_index=index,
+                dtype=self.b_dtype_name,
+                extents=self.operand_extents[index],
+                strides=self.operand_strides[index],
+            )
+
+        if c is not None:  # If we get here, we know that C is one of the operands in the problem specification.
+            c = tensor_wrapper.wrap_operand(c)
+            index = 2
+            self._check_and_set_operand(
+                c,
+                "C",
+                operand_stream_holder,
+                operand_index=index,
+                dtype=self.c_dtype_name,
+                extents=self.operand_extents[index],
+                strides=self.operand_strides[index],
+            )
+
+    def reset_operands_unchecked(
+        self,
+        *,
+        a=None,
+        b=None,
+        c=None,
+        alpha=None,
+        beta=None,
+        stream: AnyStream | int | None = None,
+    ):
+        if alpha is not None:
+            self.alpha[0] = alpha
+        if beta is not None:
+            self.beta[0] = beta
+
+        if self._memspace_eq_exespace:
+            if a is not None:
+                self._operands[0].tensor = a
+            if b is not None:
+                self._operands[1].tensor = b
+            if c is not None and self.num_operands == 3:
+                self._operands[2].tensor = c
+            return
+
+        operands = ((0, a), (1, b)) if self.num_operands == 2 else ((0, a), (1, b), (2, c))
+        _, operand_stream_holder = self._get_or_create_stream_maybe(stream)
+        for idx, value in operands:
+            if value is None:
+                continue
+            wrapped = tensor_wrapper.wrap_operand(value)
+            _realloc_if_needed_and_copy_to_mirror(
+                wrapped, self._operands[idx], self._execution_device_id, operand_stream_holder
+            )
+            self._operands_backup[idx] = wrapped
+
+    def _allocate_result_tensor_perhaps(self, *, stream_holder: utils.StreamHolder | None) -> utils.TensorHolder:
+        log_debug = self._logger.isEnabledFor(logging.DEBUG)
+        # We must handle all valid combinations of:
+        # - c-provided and c-not-provided
+        # - results in-place and out-of-place
+
+        # Create empty tensor for the result.
+        if self.num_operands == 2 or not self.options.inplace:
+            if log_debug:
+                self._logger.debug("Beginning output (empty) tensor creation...")
+                self._logger.debug(
+                    f"The output tensor shape = {self._input_traits.d_layout_traits.shape} with strides = "
+                    f"{self._input_traits.d_layout_traits.strides} and data type '{self.c_dtype_name}'."
+                )
+                self._logger.debug(
+                    f"The output tensor has batch dimensions with shape {self._batch_traits[3].shape} "
+                    f"and strides {self._batch_traits[3].strides}."
+                )
+            result = utils.create_empty_tensor(
+                self._result_class,
+                (*self._batch_traits[3].shape, *self._input_traits.d_layout_traits.shape),
+                self.c_dtype_name,
+                self._execution_device_id,
+                stream_holder,
+                # verify_strides=False because we need strides to be exactly what we
+                # request; not arbitrary if the strides aren't contiguous and dense.
+                # Otherwise, the layout parameters will mismatch what we pass to the matmul
+                # implementation.
+                verify_strides=False,
+                strides=(*self._batch_traits[3].strides, *self._input_traits.d_layout_traits.strides),
+            )
+            if log_debug:
+                self._logger.debug("The output (empty) tensor has been created.")
+        else:  # num_operands == 3 and self.options.inplace
+            result = self._operands[2]
+            self._logger.debug("The output tensor is C (in-place execution).")
+
+        if self.num_operands == 3 and not self.options.inplace:
+            result.copy_(self._operands[2], stream_holder)
+            self._logger.debug("Operand C copied to result tensor (out-of-place execution).")
+
+        return result
+
+    def _copy_result_tensor_perhaps(self, *, result: utils.TensorHolder, stream_holder: utils.StreamHolder | None) -> AnyTensor:
+        # Return the result and auxiliary outputs, if present.
+        if self._memspace_eq_exespace:
+            out = result.tensor
+        else:
+            if self.options.inplace:
+                c = self._operands_backup[2]
+                assert c is not None and c.tensor is not None, (
+                    "Internal Error. "
+                    "Inplace operation was requested, but the execution space was different from the input space, "
+                    "and we didn't keep a reference to the input tensor."
+                )
+                c.copy_(result, stream_holder=stream_holder)
+                out = c.tensor
+            else:
+                out = result.to(self._operands_device_id, stream_holder=stream_holder).tensor
+
+        return out
+
+    def execute(self, *, log_info, exec_stream_holder, operand_stream_holder) -> AnyTensor:
+        result = self._allocate_result_tensor_perhaps(stream_holder=exec_stream_holder)
+
+        a, b = self._operands[0], self._operands[1]
+        if log_info:
+            self._logger.info("Starting matrix multiplication...")
+            self._logger.info(f"{self._call_prologue}")
+
+        if self.execution.name == "cuda":
+            assert exec_stream_holder is not None
+            with utils.cuda_call_ctx(exec_stream_holder, self._blocking, timing=log_info) as (
+                _,
+                elapsed,
+            ):
+                self._function(
+                    a,
+                    b,
+                    result,
+                    self.alpha,
+                    self.beta,
+                    exec_stream_holder,
+                )
+        else:
+            with utils.host_call_ctx(timing=log_info) as elapsed:
+                self._function(
+                    a,
+                    b,
+                    result,
+                    self.alpha,
+                    self.beta,
+                    exec_stream_holder,
+                )
+
+        if log_info and elapsed.data is not None:
+            self._logger.info(f"The matrix multiplication calculation took {elapsed.data:.3f} ms to complete.")
+
+        return self._copy_result_tensor_perhaps(result=result, stream_holder=operand_stream_holder)
+
+    def release_operands(self):
+        # Same space: _operands hold direct user references.
+        # Cross space: _operands hold internal device mirrors,
+        #   _operands_backup hold direct user references.
+        # We preserve the TensorHolder wrappers and only clear the underlying tensor
+        # reference. This lets reset_operands_unchecked subsequently reuse the existing
+        # wrappers (and, in cross space, an existing device buffer) when present.
+        for op in self._operands:
+            op.tensor = None
+        if not self._memspace_eq_exespace:
+            for op in self._operands_backup:
+                op.tensor = None
+
+    def free(self):
+        # No ordering needed: CPU execution is blocking, and GPU execution operates on
+        # user-owned operands (use-after-free is the user's responsibility).
+        for attr in self.__slots__:
+            if attr != "_logger":
+                setattr(self, attr, None)
+
+
+def _effective_batch_len(value: Any) -> int | None:
+    """Effective broadcast length: None for None, len() for a non-str Sequence, else 1."""
+    if value is None:
+        return None
+    if isinstance(value, Sequence) and not isinstance(value, str):
+        return len(value)
+    return 1
+
+
+def _inplace_singleton_c_with_broadcast(a, b, c, *, alpha, beta, qualifiers) -> bool:
+    """Return True iff ``c`` is a singleton but at least one of ``a``/``b``/
+    ``qualifiers``/``alpha``/``beta`` would drive ``_zip_broadcast`` to create
+    more than one group. In that case the same ``C`` buffer is reused across
+    groups; under ``inplace=True`` the per-group updates would alias that buffer
+    and overlap, so the constructor must reject the configuration."""
+    c_is_singleton = not isinstance(c, Sequence) or len(c) == 1
+    if not c_is_singleton:
+        return False
+    return any(isinstance(x, Sequence) and len(x) > 1 for x in (a, b, qualifiers, alpha, beta))
+
+
+def _zip_broadcast(*args: Any, **kwargs: Any) -> Generator[tuple[tuple[Any, ...], dict[str, Any]], None, None]:
+    """Like zip(), but handles keyword args and broadcasts inputs to the longest length.
+
+    Raises:
+        ValueError: If the inputs are not singleton or the same length as longest input.
+
+    Returns:
+        A generator that yields tuples of (args, kwargs), broadcasted to longest length.
+    """
+    # Pre-process args such that all args are a sequence and they are all the same length or
+    # singleton.
+    arg_items: list[Sequence[Any]] = []
+    kwarg_items: dict[str, Sequence[Any]] = {}
+    lengths: list[int] = []
+
+    # Process positional arguments
+    for arg in args:
+        item = arg if isinstance(arg, Sequence) and not isinstance(arg, str) else [arg]
+        arg_items.append(item)
+        lengths.append(len(item))
+
+    # Process keyword arguments
+    for key, value in kwargs.items():
+        item = value if isinstance(value, Sequence) else [value]
+        kwarg_items[key] = item
+        lengths.append(len(item))
+
+    # lengths may be empty if no args or kwargs are provided
+    if not lengths:
+        return
+
+    # num_batch may be 0 if empty sequences are passed to args and kwargs
+    # num_batch may be positive if at least one non-empty args or kwargs are provided
+    num_batch = max(lengths)
+
+    # Check that all lengths are either 1 or num_batch; if num_batch is 0, then we already
+    # know all args will have length 0.
+    if num_batch != 0:
+        illegal_args = []
+        for i, item in enumerate(arg_items):
+            if len(item) not in (1, num_batch):
+                illegal_args.append(i)
+
+        illegal_kwargs = []
+        for key, item in kwarg_items.items():
+            if len(item) not in (1, num_batch):
+                illegal_kwargs.append(key)
+
+        if illegal_args or illegal_kwargs:
+            error_parts = []
+            if illegal_args:
+                error_parts.append(f"positional args at indices {illegal_args}")
+            if illegal_kwargs:
+                error_parts.append(f"keyword args {illegal_kwargs}")
+            raise ValueError(
+                f"All inputs must have the same length or be a singleton. "
+                f"The following do not satisfy this condition: {' and '.join(error_parts)}."
+            )
+
+    # Yield tuples of (args, kwargs) for each batch
+    for batch_idx in range(num_batch):
+        batch_args = tuple(item[batch_idx] if len(item) > 1 else item[0] for item in arg_items)
+        batch_kwargs = {key: item[batch_idx] if len(item) > 1 else item[0] for key, item in kwarg_items.items()}
+        yield (batch_args, batch_kwargs)
+
+
+class _ExplicitlyBatchedMatmul:
+    """A Matmul object that supports explicitly batched operands.
+
+    This class has the same public API as :class:`Matmul`, except that the arguments may
+    only be explicitly batched. Internally, this class just delegates to
+    :class:`_ImplicitlyBatchedMatmul` for each batch.
+
+    .. important::
+
+        Explicit batches are always defined as a Python Sequence of the same type and
+        implicit batches are always leading dimension of tensors. Otherwise, broadcasting
+        behavior would be ambiguous. For example, if some arguments are sequences of tensors
+        and other arguments are non-sequence tensors with batch dimensions.
+
+    .. seealso:: :func:`_zip_broadcast`, :class:`Matmul`, :class:`_ImplicitlyBatchedMatmul`
+
+    """
+
+    __slots__ = (
+        "_function",
+        "_logger",
+        "_original_batch_lens",
+        "batches",
+        "d_Aarray",
+        "d_Barray",
+        "d_Carray",
+        "execution",
+        "h_Aarray",
+        "h_Barray",
+        "h_Carray",
+        "last_compute_event",
+        "options",
+        "workspace_stream",
+    )
+
+    batches: Final[Sequence[_ImplicitlyBatchedMatmul]]
+    _logger: Final[LoggerLike]
+    execution: ExecutionCPU | ExecutionCUDA
+
+    def _get_or_create_stream_maybe(
+        self, stream: utils.AnyStream
+    ) -> tuple[utils.StreamHolder | None, utils.StreamHolder | None]:
+        return self.batches[0]._get_or_create_stream_maybe(stream)
+
+    def __init__(
+        self,
+        a: Sequence[AnyTensor] | AnyTensor,
+        b: Sequence[AnyTensor] | AnyTensor,
+        /,
+        c: Sequence[AnyTensor] | AnyTensor | None = None,
+        *,
+        alpha: float | complex | Sequence[float | complex] | None = None,
+        beta: float | complex | Sequence[float | complex] | None = None,
+        qualifiers: Sequence[MatrixQualifier] | MatrixQualifier | None = None,
+        options: MatmulOptions,
+        execution: ExecutionCPU | ExecutionCUDA,
+        exec_stream_holder: utils.StreamHolder | None,
+        operand_stream_holder: utils.StreamHolder | None,
+    ):
+        self.options = options
+        self._logger = self.options.logger
+        self.execution = execution
+
+        assert not (options.inplace and c is None), "Internal Error: Matmul should have rejected inplace=True with c=None."
+        assert not (
+            options.inplace and _inplace_singleton_c_with_broadcast(a, b, c, alpha=alpha, beta=beta, qualifiers=qualifiers)
+        ), "Internal Error: Matmul should have rejected inplace with broadcast c."
+
+        # None entries record "operand not provided" — distinguishable from a length-1
+        # supply, so reset_operands can reject attempts to add an operand that wasn't there
+        # originally.
+        self._original_batch_lens: dict[str, int | None] = {
+            "a": _effective_batch_len(a),
+            "b": _effective_batch_len(b),
+            "c": _effective_batch_len(c),
+            "alpha": _effective_batch_len(alpha),
+            "beta": _effective_batch_len(beta),
+        }
+
+        batches: list[_ImplicitlyBatchedMatmul] = []
+        package: str = ""
+        device_id: str | int = ""
+        dtype: str = ""
+        # Our zip functionality needs to be kwarg aware because some args are keyword only!
+        for i, (batch_args, batch_kwargs) in enumerate(
+            _zip_broadcast(
+                a,
+                b,
+                c,
+                qualifiers=qualifiers,
+                alpha=alpha,
+                beta=beta,
+            )
+        ):
+            batch = _ImplicitlyBatchedMatmul(
+                *batch_args,
+                **batch_kwargs,
+                options=options,
+                execution=execution,
+                exec_stream_holder=exec_stream_holder,
+                operand_stream_holder=operand_stream_holder,
+                _allow_mixed_batch_order=True,
+                _logger=_PrefixedLoggerAdapter(options.logger, {"prefix": f"batch {i}"}),
+            )
+            if package == "":
+                package = batch._operands_package
+            elif package != batch._operands_package:
+                raise TypeError(
+                    f"Library package mismatch: All operands must come from the same package ({batch._operands_package}) "
+                    f"as the original operand ({package})."
+                )
+            if device_id == "":
+                device_id = batch._operands_device_id
+            elif device_id != batch._operands_device_id:
+                raise TypeError(
+                    f"All operands must come from the same device ({batch._operands_device_id}) "
+                    f"as the original operand ({device_id})."
+                )
+            if dtype == "":
+                dtype = batch.a_dtype_name
+            elif dtype != batch.a_dtype_name:
+                raise ValueError(
+                    f"All batches must have the same dtype. Batch {i} has dtype "
+                    f"'{batch.a_dtype_name}', but an earlier batch has dtype '{dtype}'."
+                )
+            batches.append(batch)
+        self.batches = batches
+
+        num_pointers_all_groups = sum([max(1, batch._batch_traits[3].count) for batch in self.batches])
+        self.h_Aarray = tensor_wrapper.wrap_operand(np.empty(num_pointers_all_groups, dtype=np.int64))
+        self.h_Barray = tensor_wrapper.wrap_operand(np.empty(num_pointers_all_groups, dtype=np.int64))
+        self.h_Carray = tensor_wrapper.wrap_operand(np.empty(num_pointers_all_groups, dtype=np.int64))
+
+        if self.execution.name == "cuda":
+            # exec_stream_holder.package may be cupy/torch/etc. (not "cuda") for non-numpy
+            # operands. Safe here because NDBufferTensor.empty only reads .obj.
+            assert exec_stream_holder is not None
+            self.d_Aarray = NDBufferTensor.empty(
+                (num_pointers_all_groups,), device_id=self.execution.device_id, dtype="int64", stream_holder=exec_stream_holder
+            )
+            self.d_Barray = NDBufferTensor.empty(
+                (num_pointers_all_groups,), device_id=self.execution.device_id, dtype="int64", stream_holder=exec_stream_holder
+            )
+            self.d_Carray = NDBufferTensor.empty(
+                (num_pointers_all_groups,), device_id=self.execution.device_id, dtype="int64", stream_holder=exec_stream_holder
+            )
+            self.workspace_stream = exec_stream_holder.obj
+        else:
+            self.workspace_stream = None
+        self.last_compute_event = None
+
+    def plan(self, *, stream: AnyStream | int | None = None, log_info=False) -> None:
+        # NOTE: timing context called from Matmul.plan to avoid redundancy
+        group_batch_traits = []
+        group_mm_traits = []
+        group_qualifiers = []
+        flop_count = 0
+
+        for batch in self.batches:
+            mm_layout_checker = mm_layout_checker_getter(batch._qualifiers)
+
+            mm_traits = BLASMMTraitsView.from_input_traits(
+                batch._input_traits,
+                mm_layout_checker,
+                batch._logger,
+                lookup_table_table=CACHED_LAYOUT_CHECKERS,
+            )
+
+            batch._mm_traits = mm_traits
+
+            group_batch_traits.append(batch._batch_traits[3])
+            group_mm_traits.append(batch._mm_traits)
+            group_qualifiers.append(batch._qualifiers)
+
+            if log_info:
+                flop_count += 2 * mm_traits.M * mm_traits.N * mm_traits.K
+
+        self._function, function_name = select_blas_group_mm_function(
+            group_batch_traits,
+            group_mm_traits,
+            group_qualifiers,
+            self._logger,
+            self.execution,
+        )
+
+        if log_info:
+            self._logger.info(f"The plan found 1 suitable algorithm named {function_name}.")
+            self._logger.info(f"The base matrix multiplication FLOP count is {formatters.FLOPSStr(flop_count, 'FLOP')}.")
+
+    def reset_operands(
+        self,
+        *,
+        a=None,
+        b=None,
+        c=None,
+        alpha=None,
+        beta=None,
+        stream: AnyStream | int | None = None,
+        _operands_released=True,
+    ):
+        # An operand value of None means "skip"; otherwise its effective length must equal
+        # what was supplied at construction time, since changing the broadcast type of an
+        # operand may be unsafe.
+        unprovided: list[str] = []
+        mismatches: list[tuple[str, int, int]] = []
+        for name, value in (("a", a), ("b", b), ("c", c), ("alpha", alpha), ("beta", beta)):
+            if value is None:
+                continue
+            new_len = _effective_batch_len(value)
+            orig_len = self._original_batch_lens[name]
+            if orig_len is None:
+                unprovided.append(name)
+            elif new_len != orig_len:
+                assert new_len is not None  # value is not None implies new_len is int (mypy)
+                mismatches.append((name, orig_len, new_len))
+        if unprovided:
+            details = ", ".join(unprovided)
+            raise ValueError(
+                f"reset_operands cannot set operand(s) {details} because they were not "
+                "provided to the constructor; reset_operands may only update operands "
+                "that were originally supplied."
+            )
+        if mismatches:
+            details = ", ".join(
+                f"{name} was originally length {orig} but reset_operands received length {new}"
+                for name, orig, new in mismatches
+            )
+            raise ValueError(
+                "reset_operands received operands whose explicit-batch length does not "
+                f"match the original construction: {details}. Each operand's effective "
+                "length (len(seq) for a sequence, otherwise 1) must equal the length "
+                "supplied at construction time; broadcasting cannot be introduced or "
+                "removed by reset_operands."
+            )
+
+        # NOTE: self.batches is included in the zip, so operands are broadcast correctly
+        # when reset.
+        for batch_args, batch_kwargs in _zip_broadcast(
+            self.batches,
+            a=a,
+            b=b,
+            c=c,
+            alpha=alpha,
+            beta=beta,
+        ):
+            batch_args[0].reset_operands(
+                *batch_args[1:],
+                **batch_kwargs,
+                stream=stream,
+                _operands_released=_operands_released,
+            )
+
+    def reset_operands_unchecked(
+        self,
+        *,
+        a=None,
+        b=None,
+        c=None,
+        alpha=None,
+        beta=None,
+        stream: AnyStream | int | None = None,
+    ):
+        for batch_args, batch_kwargs in _zip_broadcast(
+            self.batches,
+            a=a,
+            b=b,
+            c=c,
+            alpha=alpha,
+            beta=beta,
+        ):
+            batch_args[0].reset_operands_unchecked(
+                *batch_args[1:],
+                **batch_kwargs,
+                stream=stream,
+            )
+
+    def execute(self, *, log_info, exec_stream_holder, operand_stream_holder) -> Sequence[AnyTensor]:
+        # Get the data type to determine which function to call
+        dtype = self.batches[0].a_dtype_name
+
+        results: list[utils.TensorHolder] = []
+
+        lo, hi = (0, 0)
+
+        for batch in self.batches:
+            # Allocate result tensors and collect their pointers
+            result = batch._allocate_result_tensor_perhaps(stream_holder=exec_stream_holder)
+            results.append(result)
+
+            # Collect all matrix pointers from each batch
+            num_pointers_this_batch = max(1, batch._batch_traits[3].count)
+            hi += num_pointers_this_batch
+            A_pointers = get_addresses_of_elements(
+                batch._batch_traits[0].shape,
+                batch._batch_traits[0].strides,
+                batch._operands[0].data_ptr,
+                itemsize=batch._operands[0].itemsize,
+            )
+            B_pointers = get_addresses_of_elements(
+                batch._batch_traits[1].shape,
+                batch._batch_traits[1].strides,
+                batch._operands[1].data_ptr,
+                itemsize=batch._operands[1].itemsize,
+            )
+            C_pointers = get_addresses_of_elements(
+                batch._batch_traits[3].shape,
+                batch._batch_traits[3].strides,
+                result.data_ptr,
+                itemsize=result.itemsize,
+            )
+
+            # Handle broadcast operands within an implicit batch
+            if len(A_pointers) < num_pointers_this_batch:
+                A_pointers = np.tile(A_pointers, num_pointers_this_batch)
+            if len(B_pointers) < num_pointers_this_batch:
+                B_pointers = np.tile(B_pointers, num_pointers_this_batch)
+
+            # Swap A and B pointers if the matrices are swapped because of layout
+            # compatibility
+            if batch._mm_traits.is_swapped_AB:
+                A_pointers, B_pointers = B_pointers, A_pointers
+
+            self.h_Aarray.tensor[lo:hi] = A_pointers
+            self.h_Barray.tensor[lo:hi] = B_pointers
+            self.h_Carray.tensor[lo:hi] = C_pointers
+
+            lo = hi
+
+        # Collect alpha and beta values
+        alpha_array = np.array([batch.alpha[0] for batch in self.batches], dtype=dtype)
+        beta_array = np.array([batch.beta[0] for batch in self.batches], dtype=dtype)
+
+        if log_info:
+            self._logger.info("Starting matrix multiplication...")
+            self._logger.info(f"{self.batches[0]._call_prologue}")
+
+        # Call the appropriate grouped batched gemm function based on dtype
+        if self.execution.name == "cpu":
+            with utils.host_call_ctx(timing=log_info) as elapsed:
+                self._function(
+                    self.h_Aarray,
+                    self.h_Barray,
+                    self.h_Carray,
+                    alpha_array,
+                    beta_array,
+                    exec_stream_holder,
+                )
+        else:
+            assert exec_stream_holder is not None
+
+            # Copy pointer arrays to device using the appropriate wrapper
+            self.d_Aarray.copy_(self.h_Aarray, exec_stream_holder)
+            self.d_Barray.copy_(self.h_Barray, exec_stream_holder)
+            self.d_Carray.copy_(self.h_Carray, exec_stream_holder)
+
+            with utils.cuda_call_ctx(exec_stream_holder, self.batches[0]._blocking, timing=log_info) as (
+                self.last_compute_event,
+                elapsed,
+            ):
+                self._function(
+                    self.d_Aarray,
+                    self.d_Barray,
+                    self.d_Carray,
+                    alpha_array,
+                    beta_array,
+                    exec_stream_holder,
+                )
+
+        if log_info and elapsed.data is not None:
+            self._logger.info(f"The matrix multiplication calculation took {elapsed.data:.3f} ms to complete.")
+
+        # Return results (copy if needed)
+        return [
+            batch._copy_result_tensor_perhaps(result=result, stream_holder=operand_stream_holder)
+            for batch, result in zip(self.batches, results, strict=True)
+        ]
+
+    def release_operands(self):
+        for b in self.batches:
+            b.release_operands()
+
+    def free(self):
+        # Ensure ordering with respect to the last computation
+        # to avoid race conditions when releasing internal resources.
+        if self.last_compute_event is not None:
+            if self.workspace_stream is not None:
+                self.workspace_stream.wait(self.last_compute_event)
+            self.last_compute_event = None
+
+        for batch in self.batches:
+            batch.free()
+
+        for attr in self.__slots__:
+            if attr != "_logger":
+                setattr(self, attr, None)
+
+
+@utils.docstring_decorator(GENERIC_MM_DOCUMENTATION, skip_missing=False)
+class Matmul:
+    """
+    Create a stateful object encapsulating the specified matrix multiplication computation
+    :math:`\\alpha a @ b + \\beta c` and the required resources to perform the operation. A
+    stateful object can be used to amortize the cost of preparation (planning in the case of
+    matrix multiplication) across multiple executions (also see the :ref:`Stateful APIs
+    <host api types>` section).
+
+    The function-form API :func:`matmul` is a convenient alternative to using stateful
+    objects for *single* use (the user needs to perform just one matrix multiplication, for
+    example), in which case there is no possibility of amortizing preparatory costs. The
+    function-form APIs are just convenience wrappers around the stateful object APIs.
+
+    Using the stateful object typically involves the following steps:
+
+    1. **Problem Specification**: Initialize the object with a defined operation and
+       options.
+    2. **Preparation**: Use :meth:`plan` to determine the best algorithmic implementation
+       for this specific matrix multiplication operation.
+    3. **Execution**: Perform the matrix multiplication computation with :meth:`execute`.
+
+    Detailed information on what's happening in the various phases described above can be
+    obtained by passing in a :class:`logging.Logger` object to :class:`MatmulOptions` or by
+    setting the appropriate options in the root logger object, which is used by default:
+
+        >>> import logging
+        >>> logging.basicConfig(
+        ...     level=logging.INFO,
+        ...     format="%(asctime)s %(levelname)-8s %(message)s",
+        ...     datefmt="%m-%d %H:%M:%S",
+        ... )
+
+    A user can select the desired logging level and, in general, take advantage of all of
+    the functionality offered by the Python `logging` module.
+
+    Args:
+        a: {a}
+
+        b: {b}
+
+        c: {c}
+
+        alpha: {alpha}
+
+        beta: {beta}
+
+        qualifiers: {qualifiers}
+
+        options: {options}
+
+        execution: {execution}
+
+        stream: {stream}
+
+    Semantics:
+        {semantics}
+
+    Thread Safety:
+        {thread_safety}
+
+    .. seealso::
+        :meth:`reset_operands`, :meth:`reset_operands_unchecked`, :meth:`execute`
+
+    Examples:
+
+        >>> import numpy as np
+        >>> import nvmath
+
+        Create two 2-D float64 ndarrays on the CPU:
+
+        >>> M, N, K = 1024, 1024, 1024
+        >>> a = np.random.rand(M, K)
+        >>> b = np.random.rand(K, N)
+
+        We will define a matrix multiplication operation using the generic matrix
+        multiplication interface.
+
+        Create a Matmul object encapsulating the problem specification above:
+
+        >>> mm = nvmath.linalg.Matmul(a, b)
+
+        Options can be provided above to control the behavior of the operation using the
+        `options` argument (see :class:`MatmulOptions`).
+
+        Next, plan the operation. The operands' layouts, qualifiers, and dtypes will be
+        considered to select an appropriate matrix multiplication:
+
+        >>> mm.plan()
+
+        Now execute the matrix multiplication, and obtain the result `r1` as a NumPy
+        ndarray.
+
+        >>> r1 = mm.execute()
+
+        Note that all :class:`Matmul` methods execute on the current stream by default.
+        Alternatively, the `stream` argument can be used to run a method on a specified
+        stream.
+
+        Let's now look at the same problem with CuPy ndarrays on the GPU.
+
+        Create a 2-D CuPy ndarray on the GPU:
+
+        >>> import cupy as cp
+        >>> a = cp.random.rand(M, K)
+        >>> b = cp.random.rand(K, N)
+
+        Create an Matmul object encapsulating the problem specification described earlier
+        and use it as a context manager.
+
+        >>> with nvmath.linalg.Matmul(a, b) as mm:
+        ...     # Plan the operation.
+        ...     mm.plan()
+        ...
+        ...     # Execute the operation to get the first result.
+        ...     r1 = mm.execute()
+        ...
+        ...     # Update operands 'a' and 'b' in-place (see reset_operands() for an
+        ...     # alternative).
+        ...     a[:] = cp.random.rand(M, K)
+        ...     b[:] = cp.random.rand(K, N)
+        ...
+        ...     # Execute the operation to get the new result.
+        ...     r2 = mm.execute()
+
+
+        All the resources used by the object are released at the end of the block.
+
+        Further examples can be found in the `examples/linalg/generic/matmul
+        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/generic/matmul>`_
+        directory.
+    """
+
+    __slots__ = (
+        "_delegate",
+        "_has_plan",
+        "_logger",
+        "_operands_released",
+        "options",
+        "valid_state",
+    )
+
+    _delegate: Final[_ImplicitlyBatchedMatmul | _ExplicitlyBatchedMatmul]
+
+    def __init__(
+        self,
+        a: AnyTensor | Sequence[AnyTensor],
+        b: AnyTensor | Sequence[AnyTensor],
+        /,
+        c: AnyTensor | Sequence[AnyTensor] | None = None,
+        *,
+        alpha: float | complex | Sequence[float | complex] | None = None,
+        beta: float | complex | Sequence[float | complex] | None = None,
+        qualifiers: MatrixQualifier | Sequence[MatrixQualifier] | None = None,
+        options: MatmulOptions | dict[str, Any] | None = None,
+        execution: ExecutionCPU | ExecutionCUDA | str | None = None,
+        stream: AnyStream | int | None = None,
+    ):
+        if isinstance(options, Sequence):
+            raise TypeError("options must be a single value.")
+        if isinstance(execution, Sequence) and not isinstance(execution, str):
+            raise TypeError("execution must be a single value.")
+        if isinstance(stream, Sequence):
+            raise TypeError("stream must be a single value.")
+
+        options = utils.check_or_create_options(MatmulOptions, options, "Matrix multiplication options")
+        assert options is not None
+        self.options = options
+        self._logger = self.options.logger
+        self._logger.info("= SPECIFICATION PHASE =")
+
+        if (
+            isinstance(a, Sequence)
+            and isinstance(b, Sequence)
+            and (c is None or isinstance(c, Sequence))
+            and (qualifiers is None or isinstance(qualifiers, Sequence))
+        ):
+            self._logger.info("The matrix multiplication is explicitly batched.")
+            delegate_class: type = _ExplicitlyBatchedMatmul
+        elif (
+            not isinstance(a, Sequence)
+            and not isinstance(b, Sequence)
+            and (c is None or not isinstance(c, Sequence))
+            and (qualifiers is None or not isinstance(qualifiers, Sequence))
+        ):
+            self._logger.info("The matrix multiplication is implicitly batched.")
+            delegate_class = _ImplicitlyBatchedMatmul
+        else:
+            raise ValueError(
+                "All arguments a, b, c, and qualifiers must be either "
+                "implicitly batched (tensors) or explicitly batched (sequences of tensors). "
+                f"a is a {type(a)}, "
+                f"b is a {type(b)}, "
+                f"c is a {type(c)}, and "
+                f"qualifiers is a {type(qualifiers)}."
+            )
+
+        empty = [
+            name
+            for name, value in (("a", a), ("b", b), ("c", c), ("qualifiers", qualifiers))
+            if isinstance(value, Sequence) and len(value) == 0
+        ]
+        if empty:
+            raise ValueError(
+                "At least one batch must be provided for explicitly batched matrix multiplication. "
+                f"The following argument(s) are empty: {', '.join(empty)}."
+            )
+
+        nested = [
+            name
+            for name, value in (("a", a), ("b", b), ("c", c), ("qualifiers", qualifiers))
+            if isinstance(value, Sequence) and isinstance(value[0], Sequence)
+        ]
+        if nested:
+            raise ValueError(
+                "Explicitly batched operands cannot be nested (sequence of sequences). "
+                f"The following argument(s) are nested: {', '.join(nested)}."
+            )
+
+        if options.inplace:
+            if c is None:
+                raise ValueError("Operation cannot be inplace if operand C is not provided.")
+            if delegate_class is _ExplicitlyBatchedMatmul and _inplace_singleton_c_with_broadcast(
+                a, b, c, alpha=alpha, beta=beta, qualifiers=qualifiers
+            ):
+                raise ValueError("Operation cannot be inplace if operand C is broadcast.")
+
+        representatives = []
+        for operand in (a, b, c):
+            if operand is None:
+                continue
+            representative = operand[0] if isinstance(operand, Sequence) else operand
+            representatives.append(tensor_wrapper.wrap_operand(representative))
+        operands_device_id = utils.get_operands_device_id(representatives)
+
+        execution = _resolve_execution(execution, operands_device_id)
+
+        self._logger.info(
+            f"The input tensors are located on device {operands_device_id}, and the execution space "
+            f"is {execution.name}, with device {getattr(execution, 'device_id', 'cpu')}."
+        )
+
+        operands_package = utils.get_operands_package(representatives)
+        internal_op_package = _resolve_internal_operand_package(execution.name, operands_package)
+        exec_stream_holder, operand_stream_holder = _get_or_create_stream_maybe(
+            execution,
+            operands_device_id,
+            operands_package,
+            internal_op_package,
+            stream,
+        )
+        holder = exec_stream_holder or operand_stream_holder
+        self._logger.info(f"The specified stream for the Matmul ctor is {holder.obj if holder else None}.")
+
+        self._delegate = delegate_class(
+            a,
+            b,
+            c=c,
+            alpha=alpha,
+            beta=beta,
+            qualifiers=qualifiers,
+            options=options,
+            execution=execution,
+            exec_stream_holder=exec_stream_holder,
+            operand_stream_holder=operand_stream_holder,
+        )
+        self._has_plan = False
+        # Track whether user-provided operands have been released
+        self._operands_released = False
+        self.valid_state = True
+        self._logger.info("The Matmul operation has been created.")
+
+    def _check_valid_matmul(self, *args, **kwargs):
+        """
+        Check if the Matmul object is alive and well.
+        """
+        if not self.valid_state:
+            raise InvalidMatmulState("The Matmul object cannot be used after resources are free'd")
+
+    def _check_planned(self, *args, **kwargs):
+        """
+        Check if plan() has been called on this Matmul object.
+        """
+        what = kwargs["what"]
+        if not self._has_plan:
+            raise RuntimeError(f"{what} cannot be performed before plan() has been called.")
+
+    def _check_valid_operands(self, *args, **kwargs):
+        """
+        Check if operands are available for operations.
+        """
+        if self._operands_released:
+            raise ValueError(
+                "Operands have been released. Call reset_operands() to provide new operands before using this method."
+            )
+
+    @utils.precondition(_check_valid_matmul)
+    @utils.precondition(_check_planned, "Execution")
+    @utils.precondition(_check_valid_operands)
+    def execute(self, *, stream: AnyStream | int | None = None) -> AnyTensor | Sequence[AnyTensor]:
+        """
+        Execute a prepared (planned) matrix multiplication.
+
+        Args:
+            stream: {stream}
+
+        Returns:
+           {result}
+        """
+        exec_stream_holder, operand_stream_holder = self._delegate._get_or_create_stream_maybe(stream)
+
+        if log_info := self._logger.isEnabledFor(logging.INFO):
+            self._logger.info("= EXECUTION PHASE =")
+            self._logger.info(
+                "The specified stream for execute() is "
+                f"{getattr(exec_stream_holder or operand_stream_holder, 'obj', 'no stream')}."
+            )
+
+        return self._delegate.execute(
+            log_info=log_info,
+            exec_stream_holder=exec_stream_holder,
+            operand_stream_holder=operand_stream_holder,
+        )
+
+    def execute_unchecked(self, *, stream: AnyStream | int | None = None) -> AnyTensor | Sequence[AnyTensor]:
+        """
+        .. experimental:: method
+
+        .. versionadded:: 1.0
+
+        This method is a performance-optimized alternative to :meth:`execute` that
+        eliminates validation and logging overhead, making it ideal for performance-critical
+        loops where operand compatibility is guaranteed by the caller.
+
+        Args:
+            stream: {stream}
+
+        Returns:
+           {result}
+
+        .. seealso::
+            :meth:`execute`: Safe, validated method for executing the matrix multiplication.
+        """
+        exec_stream_holder, operand_stream_holder = self._delegate._get_or_create_stream_maybe(stream)
+
+        return self._delegate.execute(
+            log_info=False, exec_stream_holder=exec_stream_holder, operand_stream_holder=operand_stream_holder
+        )
+
+    @utils.precondition(_check_valid_matmul)
+    def plan(self, *, stream: AnyStream | int | None = None) -> None:
+        """
+        Plan the matrix multiplication operation.
+
+        Unlike :py:meth:`nvmath.linalg.advanced.Matmul.plan`, this method takes no
+        tuning parameters. Its primary function is to find the correct matrix multiplication
+        implementation based on the operands and options provided to the constructor.
+
+        Args:
+            stream: {stream}
+
+        Returns:
+            None
+        """
+        if log_info := self._logger.isEnabledFor(logging.INFO):
+            self._logger.info("= PLANNING PHASE =")
+            if self._has_plan:
+                self._logger.info("Matmul is already planned. Replanning is skipped.")
+            else:
+                self._logger.info("Starting matrix multiplication planning...")
+
+        if self._has_plan:
+            return
+
+        with utils.host_call_ctx(timing=log_info) as elapsed:
+            self._delegate.plan(stream=stream, log_info=log_info)
+
+        if log_info and elapsed.data is not None:
+            self._logger.info(f"The matrix multiplication planning phase took {elapsed.data:.3f} ms to complete.")
+
+        self._has_plan = True
+
+    @utils.precondition(_check_valid_matmul)
+    def reset_operands(
+        self,
+        *,
+        a=None,
+        b=None,
+        c=None,
+        alpha=None,
+        beta=None,
+        stream: AnyStream | int | None = None,
+    ):
+        """
+        Reset one or more of the operands held by this :class:`Matmul` instance.
+
+        .. versionchanged:: 0.9
+            All parameters are now keyword-only.
+
+        Args:
+            a: {a}
+
+            b: {b}
+
+            c: {c}
+
+            alpha: {alpha}
+
+            beta: {beta}
+
+            stream: {stream}
+
+        Semantics:
+            - Only the operands explicitly passed are updated. At least one operand
+              is required (all of them after :meth:`release_operands`), otherwise
+              a :class:`ValueError` is raised.
+
+            - This method will perform various checks on the new operands to make sure:
+
+              - The explicit batch counts, shapes, strides, datatypes match
+                those of the old ones.
+              - The packages that the operands belong to match those of the old ones.
+              - If input tensors are on GPU, the device must match.
+
+        Examples:
+
+            >>> import cupy as cp
+            >>> import nvmath
+
+            Create two 3-D float64 ndarrays on the GPU:
+
+            >>> M, N, K = 128, 128, 256
+            >>> a = cp.random.rand(M, K)
+            >>> b = cp.random.rand(K, N)
+
+            Create an matrix multiplication object as a context manager
+
+            >>> with nvmath.linalg.Matmul(a, b) as mm:
+            ...     # Plan the operation.
+            ...     mm.plan()
+            ...
+            ...     # Execute the MM to get the first result.
+            ...     r1 = mm.execute()
+            ...
+            ...     # Reset the operands to new CuPy ndarrays.
+            ...     a_new = cp.random.rand(M, K)
+            ...     b_new = cp.random.rand(K, N)
+            ...     mm.reset_operands(a=a_new, b=b_new)
+            ...
+            ...     # Execute to get the new result corresponding to the updated operands.
+            ...     r2 = mm.execute()
+
+            Note that if only a subset of operands are reset, the operands that are not
+            reset hold their original values.
+
+            With :meth:`reset_operands`, minimal overhead is achieved as problem
+            specification and planning are only performed once.
+
+            For the particular example above, a slower alternative to calling
+            :meth:`reset_operands` would be to modify the operands in-place (e.g.,
+            ``a[:] = a_new`` and ``b[:] = b_new``), but this approach is less efficient
+            as it involves copying data rather than just updating pointers.
+
+            For more details, please refer to `inplace update example
+            <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/generic/matmul/example05_stateful_inplace.py>`_.
+        """
+        # When operands have not been released and all arguments are None, there is
+        # nothing to update, so reject the call. In the released state we fall through to
+        # the delegate, which enforces post-release completeness.
+        if not self._operands_released and all(arg is None for arg in (a, b, c, alpha, beta)):
+            raise ValueError("reset_operands() requires at least one operand to be provided.")
+
+        self._delegate.reset_operands(
+            a=a,
+            b=b,
+            c=c,
+            alpha=alpha,
+            beta=beta,
+            stream=stream,
+            # _operands_released is forwarded to the delegate so it can enforce
+            # the post-release requirement that reset_operands() must receive every
+            # operand so the operation has a fully valid state to execute on.
+            _operands_released=self._operands_released,
+        )
+        # Clear the released flag since we now have valid operands, this is
+        # guaranteed since the delegate above raised if the post-release
+        # contract was violated.
+        self._operands_released = False
+
+    def reset_operands_unchecked(
+        self,
+        *,
+        a=None,
+        b=None,
+        c=None,
+        alpha=None,
+        beta=None,
+        stream: AnyStream | int | None = None,
+    ):
+        """
+        {reset_operands_unchecked}
+        """
+        # The unchecked path waives all validation, so forwarding _operands_released
+        # is not needed because the parameter would go unused anyway.
+        # The caller is responsible for providing a complete and consistent operand set.
+        self._delegate.reset_operands_unchecked(a=a, b=b, c=c, alpha=alpha, beta=beta, stream=stream)
+        self._operands_released = False
+
+    @utils.precondition(_check_valid_matmul)
+    def release_operands(self):
+        """
+        {release_operands}
+        """
+        if self._operands_released:
+            self._logger.info("Operands have already been released; nothing to do.")
+            return
+        self._delegate.release_operands()
+        self._operands_released = True
+        self._logger.info("User-provided operands have been released.")
+
+    def free(self):
+        """Free Matmul resources.
+
+        It is recommended that the :class:`Matmul` object be used within a context,
+        but if it is not possible then this method must be called explicitly to ensure
+        that the Matmul object doesn't hold unnecessary references to operands.
+        """
+        if not self.valid_state:
+            return
+        try:
+            self._delegate.free()
+        except Exception as e:
+            self._logger.critical("Internal error: only part of the Matmul object's resources have been released.")
+            self._logger.critical(str(e))
+            raise e
+        finally:
+            self.valid_state = False
+
+        self._logger.info("The Matmul object's resources have been released.")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self.free()
+
+
+@utils.docstring_decorator(GENERIC_MM_DOCUMENTATION, skip_missing=False)
+def matmul(
+    a: AnyTensor | Sequence[AnyTensor],
+    b: AnyTensor | Sequence[AnyTensor],
+    /,
+    c: AnyTensor | Sequence[AnyTensor] | None = None,
+    *,
+    alpha: float | complex | Sequence[float | complex] | None = None,
+    beta: float | complex | Sequence[float | complex] | None = None,
+    qualifiers: MatrixQualifier | Sequence[MatrixQualifier] | None = None,
+    options: MatmulOptions | dict[str, Any] | None = None,
+    execution: ExecutionCPU | ExecutionCUDA | str | None = None,
+    stream: AnyStream | int | None = None,
+) -> AnyTensor | Sequence[AnyTensor]:
+    """
+    Perform the specified matrix multiplication computation :math:`\\alpha a @ b + \\beta
+    c`. This function-form is a wrapper around the stateful :class:`Matmul` object APIs and
+    is meant for *single* use (the user needs to perform just one matrix multiplication, for
+    example), in which case there is no possibility of amortizing preparatory costs.
+
+    Detailed information on what's happening within this function can be obtained by passing
+    in a :class:`logging.Logger` object to :class:`MatmulOptions` or by setting the
+    appropriate options in the root logger object, which is used by default:
+
+        >>> import logging
+        >>> logging.basicConfig(
+        ...     level=logging.INFO,
+        ...     format="%(asctime)s %(levelname)-8s %(message)s",
+        ...     datefmt="%m-%d %H:%M:%S",
+        ... )
+
+    A user can select the desired logging level and, in general, take advantage of all of
+    the functionality offered by the Python `logging` module.
+
+    Args:
+        a: {a}
+
+        b: {b}
+
+        c: {c}
+
+        alpha: {alpha}
+
+        beta: {beta}
+
+        qualifiers: {qualifiers}
+
+        options: {options}
+
+        execution: {execution}
+
+        stream: {stream}
+
+    Returns:
+        {result}
+
+    Semantics:
+        {semantics}
+
+    .. seealso::
+        :class:`Matmul`, :class:`MatmulOptions`, :class:`matrix_qualifiers_dtype`,
+        :class:`MatrixQualifier`
+
+    Examples:
+
+        >>> import cupy as cp
+        >>> import nvmath
+
+        Create three float32 ndarrays on the GPU:
+
+        >>> M, N, K = 128, 64, 256
+        >>> a = cp.random.rand(M, K, dtype=cp.float32)
+        >>> b = cp.random.rand(K, N, dtype=cp.float32)
+        >>> c = cp.random.rand(M, N, dtype=cp.float32)
+
+        Perform the operation :math:`\\alpha a @ b + \\beta c` using :func:`matmul`. The
+        result `r` is also a CuPy float32 ndarray:
+
+        >>> r = nvmath.linalg.matmul(a, b, c, alpha=1.23, beta=0.74)
+
+        The package current stream is used by default, but a stream can be explicitly
+        provided to the Matmul operation. This can be done if the operands are computed on a
+        different stream, for example:
+
+        >>> s = cp.cuda.Stream()
+        >>> with s:
+        ...     a = cp.random.rand(M, K)
+        ...     b = cp.random.rand(K, N)
+        >>> r = nvmath.linalg.matmul(a, b, stream=s)
+
+        The operation above runs on stream `s` and is ordered with respect to the input
+        computation.
+
+        Create NumPy ndarrays on the CPU.
+
+        >>> import numpy as np
+        >>> a = np.random.rand(M, K)
+        >>> b = np.random.rand(K, N)
+
+        Provide the NumPy ndarrays to :func:`matmul`, with the result also being a NumPy
+        ndarray:
+
+        >>> r = nvmath.linalg.matmul(a, b)
+
+    Notes:
+        - This function is a convenience wrapper around :class:`Matmul` and is
+          specifically meant for *single* use.
+
+    Further examples can be found in the `examples/linalg/generic/matmul
+    <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/generic/matmul>`_
+    directory.
+    """
+
+    with Matmul(
+        a,
+        b,
+        c=c,
+        alpha=alpha,
+        beta=beta,
+        qualifiers=qualifiers,
+        options=options,
+        execution=execution,
+        stream=stream,
+    ) as mm:
+        mm.plan(stream=stream)
+
+        r = mm.execute(stream=stream)
+
+    return r