Add diff method to XTensorVariable

Author: ricardoV94

pytensor/xtensor/type.py

@@ -499,6 +499,15 @@ def cumsum(self, dim):
     def cumprod(self, dim):
         return px.reduction.cumprod(self, dim)
 
+    def diff(self, dim, n=1):
+        """Compute the n-th discrete difference along the given dimension."""
+        slice1 = {dim: slice(1, None)}
+        slice2 = {dim: slice(None, -1)}
+        x = self
+        for _ in range(n):
+            x = x[slice1] - x[slice2]
+        return x
+
 
 class XTensorConstantSignature(tuple):
     def __eq__(self, other):

tests/xtensor/test_indexing.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 from xarray import DataArray
+from xtensor.util import xr_arange_like
 
 from pytensor.xtensor import xtensor
 from tests.xtensor.util import xr_assert_allclose, xr_function
@@ -40,3 +41,22 @@ def test_basic_indexing(labeled, indices):
     res = fn(x_test)
     expected_res = x_test[indices]
     xr_assert_allclose(res, expected_res)
+
+
+@pytest.mark.parametrize("n", ["implicit", 1, 2])
+@pytest.mark.parametrize("dim", ["a", "b"])
+def test_diff(dim, n):
+    x = xtensor(dims=("a", "b"), shape=(7, 11))
+    if n == "implicit":
+        out = x.diff(dim)
+    else:
+        out = x.diff(dim, n=n)
+
+    fn = xr_function([x], out)
+    x_test = xr_arange_like(x)
+    res = fn(x_test)
+    if n == "implicit":
+        expected_res = x_test.diff(dim)
+    else:
+        expected_res = x_test.diff(dim, n=n)
+    xr_assert_allclose(res, expected_res)