Use xr_arange_like in existing tests

Author: ricardoV94

tests/xtensor/test_shape.py

@@ -7,12 +7,16 @@
 from itertools import chain, combinations
 
 import numpy as np
-from xarray import DataArray
 from xarray import concat as xr_concat
 
 from pytensor.xtensor.shape import concat, stack
 from pytensor.xtensor.type import xtensor
-from tests.xtensor.util import xr_assert_allclose, xr_function, xr_random_like
+from tests.xtensor.util import (
+    xr_arange_like,
+    xr_assert_allclose,
+    xr_function,
+    xr_random_like,
+)
 
 
 def powerset(iterable, min_group_size=0):
@@ -42,10 +46,7 @@ def test_transpose():
     outs = [transpose(x, *perm) for perm in permutations]
 
     fn = xr_function([x], outs)
-    x_test = DataArray(
-        np.arange(np.prod(x.type.shape), dtype=x.type.dtype).reshape(x.type.shape),
-        dims=x.type.dims,
-    )
+    x_test = xr_arange_like(x)
     res = fn(x_test)
     expected_res = [x_test.transpose(*perm) for perm in permutations]
     for outs_i, res_i, expected_res_i in zip(outs, res, expected_res):
@@ -61,10 +62,7 @@ def test_stack():
     ]
 
     fn = xr_function([x], outs)
-    x_test = DataArray(
-        np.arange(np.prod(x.type.shape), dtype=x.type.dtype).reshape(x.type.shape),
-        dims=x.type.dims,
-    )
+    x_test = xr_arange_like(x)
     res = fn(x_test)
 
     expected_res = [
@@ -81,10 +79,7 @@ def test_stack_single_dim():
     assert out.type.dims == ("b", "c", "d")
 
     fn = xr_function([x], out)
-    x_test = DataArray(
-        np.arange(np.prod(x.type.shape), dtype=x.type.dtype).reshape(x.type.shape),
-        dims=x.type.dims,
-    )
+    x_test = xr_arange_like(x)
     fn.fn.dprint(print_type=True)
     res = fn(x_test)
     expected_res = x_test.stack(d=["a"])
@@ -96,10 +91,7 @@ def test_multiple_stacks():
     out = stack(x, new_dim1=("a", "b"), new_dim2=("c", "d"))
 
     fn = xr_function([x], [out])
-    x_test = DataArray(
-        np.arange(np.prod(x.type.shape), dtype=x.type.dtype).reshape(x.type.shape),
-        dims=x.type.dims,
-    )
+    x_test = xr_arange_like(x)
     res = fn(x_test)
     expected_res = x_test.stack(new_dim1=("a", "b"), new_dim2=("c", "d"))
     xr_assert_allclose(res[0], expected_res)