|
| 1 | +from ._types import Tuple, Union, array |
| 2 | +from collections.abc import Sequence |
| 3 | + |
| 4 | +def matmul(x1: array, x2: array, /) -> array: |
| 5 | + """ |
| 6 | + Computes the matrix product. |
| 7 | +
|
| 8 | + Parameters |
| 9 | + ---------- |
| 10 | + x1: array |
| 11 | + first input array. Should have a numeric data type. Must have at least one dimension. If ``x1`` is one-dimensional having shape ``(M,)`` and ``x2`` has more than one dimension, ``x1`` must be promoted to a two-dimensional array by prepending ``1`` to its dimensions (i.e., must have shape ``(1, M)``). After matrix multiplication, the prepended dimensions in the returned array must be removed. If ``x1`` has more than one dimension (including after vector-to-matrix promotion), ``shape(x1)[:-2]`` must be compatible with ``shape(x2)[:-2]`` (after vector-to-matrix promotion) (see :ref:`broadcasting`). If ``x1`` has shape ``(..., M, K)``, the innermost two dimensions form matrices on which to perform matrix multiplication. |
| 12 | + x2: array |
| 13 | + second input array. Should have a numeric data type. Must have at least one dimension. If ``x2`` is one-dimensional having shape ``(N,)`` and ``x1`` has more than one dimension, ``x2`` must be promoted to a two-dimensional array by appending ``1`` to its dimensions (i.e., must have shape ``(N, 1)``). After matrix multiplication, the appended dimensions in the returned array must be removed. If ``x2`` has more than one dimension (including after vector-to-matrix promotion), ``shape(x2)[:-2]`` must be compatible with ``shape(x1)[:-2]`` (after vector-to-matrix promotion) (see :ref:`broadcasting`). If ``x2`` has shape ``(..., K, N)``, the innermost two dimensions form matrices on which to perform matrix multiplication. |
| 14 | +
|
| 15 | + Returns |
| 16 | + ------- |
| 17 | + out: array |
| 18 | + - if both ``x1`` and ``x2`` are one-dimensional arrays having shape ``(N,)``, a zero-dimensional array containing the inner product as its only element. |
| 19 | + - if ``x1`` is a two-dimensional array having shape ``(M, K)`` and ``x2`` is a two-dimensional array having shape ``(K, N)``, a two-dimensional array containing the `conventional matrix product <https://en.wikipedia.org/wiki/Matrix_multiplication>`_ and having shape ``(M, N)``. |
| 20 | + - if ``x1`` is a one-dimensional array having shape ``(K,)`` and ``x2`` is an array having shape ``(..., K, N)``, an array having shape ``(..., N)`` (i.e., prepended dimensions during vector-to-matrix promotion must be removed) and containing the `conventional matrix product <https://en.wikipedia.org/wiki/Matrix_multiplication>`_. |
| 21 | + - if ``x1`` is an array having shape ``(..., M, K)`` and ``x2`` is a one-dimensional array having shape ``(K,)``, an array having shape ``(..., M)`` (i.e., appended dimensions during vector-to-matrix promotion must be removed) and containing the `conventional matrix product <https://en.wikipedia.org/wiki/Matrix_multiplication>`_. |
| 22 | + - if ``x1`` is a two-dimensional array having shape ``(M, K)`` and ``x2`` is an array having shape ``(..., K, N)``, an array having shape ``(..., M, N)`` and containing the `conventional matrix product <https://en.wikipedia.org/wiki/Matrix_multiplication>`_ for each stacked matrix. |
| 23 | + - if ``x1`` is an array having shape ``(..., M, K)`` and ``x2`` is a two-dimensional array having shape ``(K, N)``, an array having shape ``(..., M, N)`` and containing the `conventional matrix product <https://en.wikipedia.org/wiki/Matrix_multiplication>`_ for each stacked matrix. |
| 24 | + - if either ``x1`` or ``x2`` has more than two dimensions, an array having a shape determined by :ref:`broadcasting` ``shape(x1)[:-2]`` against ``shape(x2)[:-2]`` and containing the `conventional matrix product <https://en.wikipedia.org/wiki/Matrix_multiplication>`_ for each stacked matrix. |
| 25 | +
|
| 26 | + The returned array must have a data type determined by :ref:`type-promotion`. |
| 27 | +
|
| 28 | + Notes |
| 29 | + ----- |
| 30 | + - The ``matmul`` function must implement the same semantics as the built-in ``@`` operator (see `PEP 465 <https://www.python.org/dev/peps/pep-0465>`_). |
| 31 | +
|
| 32 | + **Raises** |
| 33 | +
|
| 34 | + - if either ``x1`` or ``x2`` is a zero-dimensional array. |
| 35 | + - if ``x1`` is a one-dimensional array having shape ``(K,)``, ``x2`` is a one-dimensional array having shape ``(L,)``, and ``K != L``. |
| 36 | + - if ``x1`` is a one-dimensional array having shape ``(K,)``, ``x2`` is an array having shape ``(..., L, N)``, and ``K != L``. |
| 37 | + - if ``x1`` is an array having shape ``(..., M, K)``, ``x2`` is a one-dimensional array having shape ``(L,)``, and ``K != L``. |
| 38 | + - if ``x1`` is an array having shape ``(..., M, K)``, ``x2`` is an array having shape ``(..., L, N)``, and ``K != L``. |
| 39 | + """ |
| 40 | + |
| 41 | +def matrix_transpose(x: array, /) -> array: |
| 42 | + """ |
| 43 | + Transposes a matrix (or a stack of matrices) ``x``. |
| 44 | +
|
| 45 | + Parameters |
| 46 | + ---------- |
| 47 | + x: array |
| 48 | + input array having shape ``(..., M, N)`` and whose innermost two dimensions form ``MxN`` matrices. |
| 49 | +
|
| 50 | + Returns |
| 51 | + ------- |
| 52 | + out: array |
| 53 | + an array containing the transpose for each matrix and having shape ``(..., N, M)``. The returned array must have the same data type as ``x``. |
| 54 | + """ |
| 55 | + |
| 56 | +def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> array: |
| 57 | + """ |
| 58 | + Returns a tensor contraction of ``x1`` and ``x2`` over specific axes. |
| 59 | +
|
| 60 | + Parameters |
| 61 | + ---------- |
| 62 | + x1: array |
| 63 | + first input array. Should have a numeric data type. |
| 64 | + x2: array |
| 65 | + second input array. Must be compatible with ``x1`` for all non-contracted axes (see :ref:`broadcasting`). Should have a numeric data type. |
| 66 | + axes: Union[int, Tuple[Sequence[int], Sequence[int]]] |
| 67 | + number of axes (dimensions) to contract or explicit sequences of axes (dimensions) for ``x1`` and ``x2``, respectively. |
| 68 | +
|
| 69 | + If ``axes`` is an ``int`` equal to ``N``, then contraction must be performed over the last ``N`` axes of ``x1`` and the first ``N`` axes of ``x2`` in order. The size of each corresponding axis (dimension) must match. Must be nonnegative. |
| 70 | +
|
| 71 | + - If ``N`` equals ``0``, the result is the tensor (outer) product. |
| 72 | + - If ``N`` equals ``1``, the result is the tensor dot product. |
| 73 | + - If ``N`` equals ``2``, the result is the tensor double contraction (default). |
| 74 | +
|
| 75 | + If ``axes`` is a tuple of two sequences ``(x1_axes, x2_axes)``, the first sequence must apply to ``x`` and the second sequence to ``x2``. Both sequences must have the same length. Each axis (dimension) ``x1_axes[i]`` for ``x1`` must have the same size as the respective axis (dimension) ``x2_axes[i]`` for ``x2``. Each sequence must consist of unique (nonnegative) integers that specify valid axes for each respective array. |
| 76 | +
|
| 77 | + Returns |
| 78 | + ------- |
| 79 | + out: array |
| 80 | + an array containing the tensor contraction whose shape consists of the non-contracted axes (dimensions) of the first array ``x1``, followed by the non-contracted axes (dimensions) of the second array ``x2``. The returned array must have a data type determined by :ref:`type-promotion`. |
| 81 | +
|
| 82 | + Notes |
| 83 | + ----- |
| 84 | + - Contracted axes (dimensions) must not be broadcasted. |
| 85 | + """ |
| 86 | + |
| 87 | +def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array: |
| 88 | + """ |
| 89 | + Computes the (vector) dot product of two arrays. |
| 90 | +
|
| 91 | + Parameters |
| 92 | + ---------- |
| 93 | + x1: array |
| 94 | + first input array. Should have a numeric data type. |
| 95 | + x2: array |
| 96 | + second input array. Must be compatible with ``x1`` (see :ref:`broadcasting`). Should have a numeric data type. |
| 97 | + axis:int |
| 98 | + axis over which to compute the dot product. Must be an integer on the interval ``[-N, N)``, where ``N`` is the rank (number of dimensions) of the shape determined according to :ref:`broadcasting`. If specified as a negative integer, the function must determine the axis along which to compute the dot product by counting backward from the last dimension (where ``-1`` refers to the last dimension). By default, the function must compute the dot product over the last axis. Default: ``-1``. |
| 99 | +
|
| 100 | + Returns |
| 101 | + ------- |
| 102 | + out: array |
| 103 | + if ``x1`` and ``x2`` are both one-dimensional arrays, a zero-dimensional containing the dot product; otherwise, a non-zero-dimensional array containing the dot products and having rank ``N-1``, where ``N`` is the rank (number of dimensions) of the shape determined according to :ref:`broadcasting`. The returned array must have a data type determined by :ref:`type-promotion`. |
| 104 | +
|
| 105 | +
|
| 106 | + **Raises** |
| 107 | +
|
| 108 | + - if provided an invalid ``axis``. |
| 109 | + - if the size of the axis over which to compute the dot product is not the same for both ``x1`` and ``x2``. |
| 110 | + """ |
| 111 | + |
| 112 | +__all__ = ['matmul', 'matrix_transpose', 'tensordot', 'vecdot'] |
0 commit comments