Add 2023.12 axis restrictions to vecdot() and cross()

Author: asmeurer

array_api_strict/_linear_algebra_functions.py

@@ -8,8 +8,8 @@
 from __future__ import annotations
 
 from ._dtypes import _numeric_dtypes
-
 from ._array_object import Array
+from ._flags import get_array_api_strict_flags
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
@@ -54,6 +54,19 @@ def matrix_transpose(x: Array, /) -> Array:
 def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array:
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
         raise TypeError('Only numeric dtypes are allowed in vecdot')
+
+    if get_array_api_strict_flags()['api_version'] >= '2023.12':
+        if axis >= 0:
+            raise ValueError("axis must be negative in vecdot")
+        elif axis < min(-1, -x1.ndim, -x2.ndim):
+            raise ValueError("axis is out of bounds for x1 and x2")
+
+    # In versions if the standard prior to 2023.12, vecdot applied axis after
+    # broadcasting. This is different from applying it before broadcasting
+    # when axis is nonnegative. The below code keeps this behavior for
+    # 2022.12, primarily for backwards compatibility. Note that the behavior
+    # is unambiguous when axis is negative, so the below code should work
+    # correctly in that case regardless of which version is used.
     ndim = max(x1.ndim, x2.ndim)
     x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape)
     x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape)

array_api_strict/linalg.py

@@ -80,6 +80,17 @@ def cross(x1: Array, x2: Array, /, *, axis: int = -1) -> Array:
     # Note: this is different from np.cross(), which allows dimension 2
     if x1.shape[axis] != 3:
         raise ValueError('cross() dimension must equal 3')
+
+    if get_array_api_strict_flags()['api_version'] >= '2023.12':
+        if axis >= 0:
+            raise ValueError("axis must be negative in cross")
+        elif axis < min(-1, -x1.ndim, -x2.ndim):
+            raise ValueError("axis is out of bounds for x1 and x2")
+
+        # Prior to 2023.12, there was ambiguity in the standard about whether
+        # positive axis applied before or after broadcasting. NumPy applies
+        # the axis before broadcasting. Since that behavior is what has always
+        # been implemented here, we keep it for backwards compatibility.
     return Array._new(np.cross(x1._array, x2._array, axis=axis))
 
 @requires_extension('linalg')

array_api_strict/tests/test_linalg.py

@@ -0,0 +1,133 @@
+import pytest
+
+from .._flags import set_array_api_strict_flags
+
+import array_api_strict as xp
+
+# TODO: Maybe all of these exceptions should be IndexError?
+
+# Technically this is linear_algebra, not linalg, but it's simpler to keep
+# both of these tests together
+def test_vecdot_2023_12():
+    # Test the axis < 0 restriction for 2023.12, and also the 2022.12 axis >=
+    # 0 behavior (which is primarily kept for backwards compatibility).
+
+    a = xp.ones((2, 3, 4, 5))
+    b = xp.ones((   3, 4, 1))
+
+    # 2022.12 behavior, which is to apply axis >= 0 after broadcasting
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=0))
+    assert xp.linalg.vecdot(a, b, axis=1).shape == (2, 4, 5)
+    assert xp.linalg.vecdot(a, b, axis=2).shape == (2, 3, 5)
+    # This is disallowed because the arrays must have the same values before
+    # broadcasting
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=-1))
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=-4))
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=3))
+
+    # Out-of-bounds axes even after broadcasting
+    pytest.raises(IndexError, lambda: xp.linalg.vecdot(a, b, axis=4))
+    pytest.raises(IndexError, lambda: xp.linalg.vecdot(a, b, axis=-5))
+
+    # negative axis behavior is unambiguous when it's within the bounds of
+    # both arrays before broadcasting
+    assert xp.linalg.vecdot(a, b, axis=-2).shape == (2, 3, 5)
+    assert xp.linalg.vecdot(a, b, axis=-3).shape == (2, 4, 5)
+
+    # 2023.12 behavior, which is to only allow axis < 0 and axis >=
+    # min(x1.ndim, x2.ndim), which is unambiguous
+    with pytest.warns(UserWarning):
+        set_array_api_strict_flags(api_version='2023.12')
+
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=0))
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=1))
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=2))
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=3))
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=-1))
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=-4))
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=4))
+    pytest.raises(ValueError, lambda: xp.linalg.vecdot(a, b, axis=-5))
+
+    assert xp.linalg.vecdot(a, b, axis=-2).shape == (2, 3, 5)
+    assert xp.linalg.vecdot(a, b, axis=-3).shape == (2, 4, 5)
+
+@pytest.mark.parametrize('api_version', ['2021.12', '2022.12', '2023.12'])
+def test_cross(api_version):
+    # This test tests everything that should be the same across all supported
+    # API versions.
+
+    if api_version != '2022.12':
+        with pytest.warns(UserWarning):
+            set_array_api_strict_flags(api_version=api_version)
+    else:
+        set_array_api_strict_flags(api_version=api_version)
+
+    a = xp.ones((2, 4, 5, 3))
+    b = xp.ones((   4, 1, 3))
+    assert xp.linalg.cross(a, b, axis=-1).shape == (2, 4, 5, 3)
+
+    a = xp.ones((2, 4, 3, 5))
+    b = xp.ones((   4, 3, 1))
+    assert xp.linalg.cross(a, b, axis=-2).shape == (2, 4, 3, 5)
+
+    # This is disallowed because the axes must equal 3 before broadcasting
+    a = xp.ones((3, 2, 3, 5))
+    b = xp.ones((   2, 1, 1))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=-1))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=-2))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=-3))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=-4))
+
+    # Out-of-bounds axes even after broadcasting
+    pytest.raises(IndexError, lambda: xp.linalg.cross(a, b, axis=4))
+    pytest.raises(IndexError, lambda: xp.linalg.cross(a, b, axis=-5))
+
+@pytest.mark.parametrize('api_version', ['2021.12', '2022.12'])
+def test_cross_2022_12(api_version):
+    # Test the 2022.12 axis >= 0 behavior, which is primarily kept for
+    # backwards compatibility. Note that unlike vecdot, array_api_strict
+    # cross() never implemented the "after broadcasting" axis behavior, but
+    # just reused NumPy cross(), which applies axes before broadcasting.
+    if api_version != '2022.12':
+        with pytest.warns(UserWarning):
+            set_array_api_strict_flags(api_version=api_version)
+    else:
+        set_array_api_strict_flags(api_version=api_version)
+
+    a = xp.ones((3, 2, 4, 5))
+    b = xp.ones((3, 2, 4, 1))
+    assert xp.linalg.cross(a, b, axis=0).shape == (3, 2, 4, 5)
+
+    # ambiguous case
+    a = xp.ones((   3, 4, 5))
+    b = xp.ones((3, 2, 4, 1))
+    assert xp.linalg.cross(a, b, axis=0).shape == (3, 2, 4, 5)
+
+def test_cross_2023_12():
+    # 2023.12 behavior, which is to only allow axis < 0 and axis >=
+    # min(x1.ndim, x2.ndim), which is unambiguous
+    with pytest.warns(UserWarning):
+        set_array_api_strict_flags(api_version='2023.12')
+
+    a = xp.ones((3, 2, 4, 5))
+    b = xp.ones((3, 2, 4, 1))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=0))
+
+    a = xp.ones((   3, 4, 5))
+    b = xp.ones((3, 2, 4, 1))
+    pytest.raises(ValueError, lambda: xp. linalg.cross(a, b, axis=0))
+
+    a = xp.ones((2, 4, 5, 3))
+    b = xp.ones((   4, 1, 3))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=0))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=1))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=2))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=3))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=-2))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=-3))
+    pytest.raises(ValueError, lambda: xp.linalg.cross(a, b, axis=-4))
+
+    pytest.raises(IndexError, lambda: xp.linalg.cross(a, b, axis=4))
+    pytest.raises(IndexError, lambda: xp.linalg.cross(a, b, axis=-5))
+
+    assert xp.linalg.cross(a, b, axis=-1).shape == (2, 4, 5, 3)

array_api_strict/tests/test_statistical_functions.py

@@ -1,5 +1,7 @@
 import pytest
 
+from .._flags import set_array_api_strict_flags
+
 import array_api_strict as xp
 
 @pytest.mark.parametrize('func_name', ['sum', 'prod', 'trace'])
@@ -20,7 +22,7 @@ def test_sum_prod_trace_2023_12(func_name):
     assert func(a_int).dtype == xp.int64
 
     with pytest.warns(UserWarning):
-        xp.set_array_api_strict_flags(api_version='2023.12')
+        set_array_api_strict_flags(api_version='2023.12')
 
     assert func(a_real).dtype == xp.float32
     assert func(a_complex).dtype == xp.complex64