torch: add diff

Author: ev-br

array_api_compat/torch/_aliases.py

@@ -8,6 +8,7 @@
                                clip as _aliases_clip,
                                unstack as _aliases_unstack,
                                cumulative_sum as _aliases_cumulative_sum,
+                               cumulative_prod as _aliases_cumulative_prod,
                                )
 from .._internal import get_xp
 
@@ -499,6 +500,20 @@ def nonzero(x: array, /, **kwargs) -> Tuple[array, ...]:
         raise ValueError("nonzero() does not support zero-dimensional arrays")
     return torch.nonzero(x, as_tuple=True, **kwargs)
 
+
+# torch uses `dim` instead of `axis`
+def diff(
+    x: array,
+    /,
+    *,
+    axis: int = -1,
+    n: int = 1,
+    prepend: Optional[array] = None,
+    append: Optional[array] = None,
+) -> array:
+    return torch.diff(x, dim=axis, n=n, prepend=prepend, append=append)
+
+
 # torch uses `dim` instead of `axis`
 def count_nonzero(
     x: array,
@@ -748,7 +763,7 @@ def sign(x: array, /) -> array:
            'permute_dims', 'bitwise_invert', 'newaxis', 'conj', 'add',
            'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
            'bitwise_right_shift', 'bitwise_xor', 'copysign', 'count_nonzero',
-           'divide',
+           'diff', 'divide',
            'equal', 'floor_divide', 'greater', 'greater_equal', 'hypot',
            'less', 'less_equal', 'logaddexp', 'maximum', 'minimum',
            'multiply', 'not_equal', 'pow', 'remainder', 'subtract', 'max',