Add diff method to XTensorVariable

Author: ricardoV94

pytensor/xtensor/type.py

@@ -502,6 +502,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

@@ -117,3 +117,22 @@ def test_single_vector_indexing_interacting_with_exisiting_dim():
     res = fn(x_test, xidx_test)
     expected_res = x_test[xidx_test.rename(a="b"), 1:]
     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)