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)