Add preliminary diff() function for 2024.12

Author: asmeurer

array_api_strict/__init__.py

@@ -305,9 +305,9 @@
 
 __all__ += ["cumulative_sum", "max", "mean", "min", "prod", "std", "sum", "var"]
 
-from ._utility_functions import all, any
+from ._utility_functions import all, any, diff
 
-__all__ += ["all", "any"]
+__all__ += ["all", "any", "diff"]
 
 from ._array_object import Device
 __all__ += ["Device"]

array_api_strict/_utility_functions.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from ._array_object import Array
+from ._flags import requires_api_version
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
@@ -37,3 +38,25 @@ def any(
     See its docstring for more information.
     """
     return Array._new(np.asarray(np.any(x._array, axis=axis, keepdims=keepdims)), device=x.device)
+
+@requires_api_version('2024.12')
+def diff(
+    x: Array,
+    /,
+    *,
+    axis: int = -1,
+    n: int = 1,
+    prepend: Optional[Array] = None,
+    append: Optional[Array] = None,
+) -> Array:
+    # NumPy does not support prepend=None or append=None
+    kwargs = dict(axis=axis, n=n)
+    if prepend is not None:
+        if prepend.device != x.device:
+            raise ValueError(f"Arrays from two different devices ({prepend.device} and {x.device}) can not be combined.")
+        kwargs['prepend'] = prepend._array
+    if append is not None:
+        if append.device != x.device:
+            raise ValueError(f"Arrays from two different devices ({append.device} and {x.device}) can not be combined.")
+        kwargs['append'] = append._array
+    return Array._new(np.diff(x._array, **kwargs), device=x.device)

array_api_strict/tests/test_flags.py

@@ -282,6 +282,10 @@ def test_fft(func_name):
     'unstack': lambda: xp.unstack(xp.ones((3, 3)), axis=0),
 }
 
+api_version_2024_12_examples = {
+    'diff': lambda: xp.diff(xp.asarray([0, 1, 2])),
+}
+
 @pytest.mark.parametrize('func_name', api_version_2023_12_examples.keys())
 def test_api_version_2023_12(func_name):
     func = api_version_2023_12_examples[func_name]
@@ -300,6 +304,28 @@ def test_api_version_2023_12(func_name):
     set_array_api_strict_flags(api_version='2022.12')
     pytest.raises(RuntimeError, func)
 
+@pytest.mark.parametrize('func_name', api_version_2024_12_examples.keys())
+def test_api_version_2024_12(func_name):
+    func = api_version_2024_12_examples[func_name]
+
+    # By default, these functions should error
+    pytest.raises(RuntimeError, func)
+
+    # In 2022.12 and 2023.12, these functions should error
+    set_array_api_strict_flags(api_version='2022.12')
+    pytest.raises(RuntimeError, func)
+    set_array_api_strict_flags(api_version='2023.12')
+    pytest.raises(RuntimeError, func)
+
+    # They should not error in 2024.12
+    with pytest.warns(UserWarning):
+        set_array_api_strict_flags(api_version='2024.12')
+    func()
+
+    # Test the behavior gets updated properly
+    set_array_api_strict_flags(api_version='2023.12')
+    pytest.raises(RuntimeError, func)
+
 def test_disabled_extensions():
     # Test that xp.extension errors when an extension is disabled, and that
     # xp.__all__ is updated properly.