Add tests for disabling linalg and fft extensions

Author: asmeurer

array_api_strict/tests/test_flags.py

@@ -4,6 +4,8 @@
 from .. import (asarray, unique_all, unique_counts, unique_inverse,
                 unique_values, nonzero)
 
+import array_api_strict as xp
+
 import pytest
 
 @pytest.fixture(autouse=True)
@@ -76,3 +78,92 @@ def test_data_dependent_shapes():
     pytest.raises(RuntimeError, lambda: unique_values(a))
     pytest.raises(RuntimeError, lambda: nonzero(a))
     pytest.raises(RuntimeError, lambda: a[mask])
+
+linalg_examples = {
+    'cholesky': lambda: xp.linalg.cholesky(xp.eye(3)),
+    'cross': lambda: xp.linalg.cross(xp.asarray([1, 0, 0]), xp.asarray([0, 1, 0])),
+    'det': lambda: xp.linalg.det(xp.eye(3)),
+    'diagonal': lambda: xp.linalg.diagonal(xp.eye(3)),
+    'eigh': lambda: xp.linalg.eigh(xp.eye(3)),
+    'eigvalsh': lambda: xp.linalg.eigvalsh(xp.eye(3)),
+    'inv': lambda: xp.linalg.inv(xp.eye(3)),
+    'matmul': lambda: xp.linalg.matmul(xp.eye(3), xp.eye(3)),
+    'matrix_norm': lambda: xp.linalg.matrix_norm(xp.eye(3)),
+    'matrix_power': lambda: xp.linalg.matrix_power(xp.eye(3), 2),
+    'matrix_rank': lambda: xp.linalg.matrix_rank(xp.eye(3)),
+    'matrix_transpose': lambda: xp.linalg.matrix_transpose(xp.eye(3)),
+    'outer': lambda: xp.linalg.outer(xp.asarray([1, 2, 3]), xp.asarray([4, 5, 6])),
+    'pinv': lambda: xp.linalg.pinv(xp.eye(3)),
+    'qr': lambda: xp.linalg.qr(xp.eye(3)),
+    'slogdet': lambda: xp.linalg.slogdet(xp.eye(3)),
+    'solve': lambda: xp.linalg.solve(xp.eye(3), xp.eye(3)),
+    'svd': lambda: xp.linalg.svd(xp.eye(3)),
+    'svdvals': lambda: xp.linalg.svdvals(xp.eye(3)),
+    'tensordot': lambda: xp.linalg.tensordot(xp.eye(3), xp.eye(3)),
+    'trace': lambda: xp.linalg.trace(xp.eye(3)),
+    'vecdot': lambda: xp.linalg.vecdot(xp.asarray([1, 2, 3]), xp.asarray([4, 5, 6])),
+    'vector_norm': lambda: xp.linalg.vector_norm(xp.asarray([1., 2., 3.])),
+}
+
+assert set(linalg_examples) == set(xp.linalg.__all__)
+
+linalg_main_namespace_examples = {
+    'matmul': lambda: xp.matmul(xp.eye(3), xp.eye(3)),
+    'matrix_transpose': lambda: xp.matrix_transpose(xp.eye(3)),
+    'tensordot': lambda: xp.tensordot(xp.eye(3), xp.eye(3)),
+    'vecdot': lambda: xp.vecdot(xp.asarray([1, 2, 3]), xp.asarray([4, 5, 6])),
+}
+
+assert set(linalg_main_namespace_examples) == set(xp.__all__) & set(xp.linalg.__all__)
+
+@pytest.mark.parametrize('func_name', linalg_examples.keys())
+def test_linalg(func_name):
+    func = linalg_examples[func_name]
+    if func_name in linalg_main_namespace_examples:
+        main_namespace_func = linalg_main_namespace_examples[func_name]
+    else:
+        main_namespace_func = lambda: None
+
+    # First make sure the example actually works
+    func()
+    main_namespace_func()
+
+    set_array_api_strict_flags(enabled_extensions=())
+    pytest.raises(RuntimeError, func)
+    main_namespace_func()
+
+    set_array_api_strict_flags(enabled_extensions=('linalg',))
+    func()
+    main_namespace_func()
+
+fft_examples = {
+    'fft': lambda: xp.fft.fft(xp.asarray([0j, 1j, 0j, 0j])),
+    'ifft': lambda: xp.fft.ifft(xp.asarray([0j, 1j, 0j, 0j])),
+    'fftn': lambda: xp.fft.fftn(xp.asarray([[0j, 1j], [0j, 0j]])),
+    'ifftn': lambda: xp.fft.ifftn(xp.asarray([[0j, 1j], [0j, 0j]])),
+    'rfft': lambda: xp.fft.rfft(xp.asarray([0., 1., 0., 0.])),
+    'irfft': lambda: xp.fft.irfft(xp.asarray([0j, 1j, 0j, 0j])),
+    'rfftn': lambda: xp.fft.rfftn(xp.asarray([[0., 1.], [0., 0.]])),
+    'irfftn': lambda: xp.fft.irfftn(xp.asarray([[0j, 1j], [0j, 0j]])),
+    'hfft': lambda: xp.fft.hfft(xp.asarray([0j, 1j, 0j, 0j])),
+    'ihfft': lambda: xp.fft.ihfft(xp.asarray([0., 1., 0., 0.])),
+    'fftfreq': lambda: xp.fft.fftfreq(4),
+    'rfftfreq': lambda: xp.fft.rfftfreq(4),
+    'fftshift': lambda: xp.fft.fftshift(xp.asarray([0j, 1j, 0j, 0j])),
+    'ifftshift': lambda: xp.fft.ifftshift(xp.asarray([0j, 1j, 0j, 0j])),
+}
+
+assert set(fft_examples) == set(xp.fft.__all__)
+
+@pytest.mark.parametrize('func_name', fft_examples.keys())
+def test_fft(func_name):
+    func = fft_examples[func_name]
+
+    # First make sure the example actually works
+    func()
+
+    set_array_api_strict_flags(enabled_extensions=())
+    pytest.raises(RuntimeError, func)
+
+    set_array_api_strict_flags(enabled_extensions=('fft',))
+    func()