Allow negative axes in tensordot

asmeurer · kgryte · web-flow · commit 8c42f02add07 · 2023-09-19T17:29:13.000-07:00
PR-URL: #625 Co-authored-by: Athan Reines <kgryte@gmail.com> Reviewed-by: Athan Reines <kgryte@gmail.com>
diff --git a/src/array_api_stubs/_draft/linear_algebra_functions.py b/src/array_api_stubs/_draft/linear_algebra_functions.py
@@ -92,15 +92,15 @@ def tensordot(
            Contracted axes (dimensions) must not be broadcasted.
 
     axes: Union[int, Tuple[Sequence[int], Sequence[int]]]
-        number of axes (dimensions) to contract or explicit sequences of axes (dimensions) for ``x1`` and ``x2``, respectively.
+        number of axes (dimensions) to contract or explicit sequences of axis (dimension) indices for ``x1`` and ``x2``, respectively.
 
         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.
 
         -   If ``N`` equals ``0``, the result is the tensor (outer) product.
         -   If ``N`` equals ``1``, the result is the tensor dot product.
         -   If ``N`` equals ``2``, the result is the tensor double contraction (default).
 
-        If ``axes`` is a tuple of two sequences ``(x1_axes, x2_axes)``, the first sequence must apply to ``x1`` 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.
+        If ``axes`` is a tuple of two sequences ``(x1_axes, x2_axes)``, the first sequence must apply to ``x1`` 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 index referred to in a sequence must be unique. If ``x1`` has rank (i.e, number of dimensions) ``N``, a valid ``x1`` axis must reside on the half-open interval ``[-N, N)``. If ``x2`` has rank ``M``, a valid ``x2`` axis must reside on the half-open interval ``[-M, M)``.
 
 
     .. note::
