Refactor stacking functions, add dstack and column_stack

pytensor/tensor/basic.py

@@ -2758,15 +2758,34 @@ def concatenate(tensor_list, axis=0):
     return join(axis, *tensor_list)
 
 
-def horizontal_stack(*args):
+def hstack(tup):
     r"""Stack arrays in sequence horizontally (column wise)."""
-    # Note: 'horizontal_stack' and 'vertical_stack' do not behave exactly like
-    # Numpy's hstack and vstack functions. This is intended, because Numpy's
-    # functions have potentially confusing/incoherent behavior (try them on 1D
-    # arrays). If this is fixed in a future version of Numpy, it may be worth
-    # trying to get closer to Numpy's way of doing things. In the meantime,
-    # better keep different names to emphasize the implementation divergences.
 
+    arrs = atleast_1d(*tup)
+    if not isinstance(arrs, list):
+        arrs = [arrs]
+    # As a special case, dimension 0 of 1-dimensional arrays is "horizontal"
+    if arrs and arrs[0].ndim == 1:
+        return concatenate(arrs, axis=0)
+    else:
+        return concatenate(arrs, axis=1)
+
+
+def vstack(tup):
+    r"""Stack arrays in sequence vertically (row wise)."""
+
+    arrs = atleast_2d(*tup)
+    if not isinstance(arrs, list):
+        arrs = [arrs]
+
+    return concatenate(arrs, axis=0)
+
+
+def horizontal_stack(*args):
+    warnings.warn(
+        "horizontal_stack was renamed to hstack and will be removed in a future release",
+        FutureWarning,
+    )
     if len(args) < 2:
         raise ValueError("Too few arguments")
 
@@ -2781,8 +2800,10 @@ def horizontal_stack(*args):
 
 
 def vertical_stack(*args):
-    r"""Stack arrays in sequence vertically (row wise)."""
-
+    warnings.warn(
+        "vertical_stack was renamed to vstack and will be removed in a future release",
+        FutureWarning,
+    )
     if len(args) < 2:
         raise ValueError("Too few arguments")
 
@@ -2796,6 +2817,33 @@ def vertical_stack(*args):
     return concatenate(_args, axis=0)
 
 
+def dstack(tup):
+    r"""Stack arrays in sequence along third axis (depth wise)."""
+
+    # arrs = atleast_3d(*tup, left=False)
+    # if not isinstance(arrs, list):
+    #     arrs = [arrs]
+    arrs = []
+    for arr in tup:
+        if arr.ndim == 1:
+            arr = atleast_3d([arr], left=False)
+        else:
+            arr = atleast_3d(arr, left=False)
+        arrs.append(arr)
+    return concatenate(arrs, 2)
+
+
+def column_stack(tup):
+    r"""Stack 1-D arrays as columns into a 2-D array."""
+
+    arrays = []
+    for arr in tup:
+        if arr.ndim < 2:
+            arr = atleast_2d(arr).transpose()
+        arrays.append(arr)
+    return concatenate(arrays, 1)
+
+
 def is_flat(var, ndim=1):
     """
     Verifies the dimensionality of the var is equal to
@@ -4298,8 +4346,12 @@ def ix_(*args):
     "tile",
     "flatten",
     "is_flat",
+    "vstack",
+    "hstack",
     "vertical_stack",
     "horizontal_stack",
+    "dstack",
+    "column_stack",
     "get_vector_length",
     "concatenate",
     "stack",

tests/tensor/test_basic.py

@@ -41,9 +41,11 @@
     atleast_Nd,
     cast,
     choose,
+    column_stack,
     constant,
     default,
     diag,
+    dstack,
     expand_dims,
     extract_constant,
     eye,
@@ -55,6 +57,7 @@
     get_underlying_scalar_constant_value,
     get_vector_length,
     horizontal_stack,
+    hstack,
     identity_like,
     infer_static_shape,
     inverse_permutation,
@@ -86,6 +89,7 @@
     triu_indices,
     triu_indices_from,
     vertical_stack,
+    vstack,
     zeros_like,
 )
 from pytensor.tensor.blockwise import Blockwise
@@ -1667,11 +1671,25 @@ def test_join_matrix_ints(self):
         assert (np.asarray(grad(s.sum(), b).eval()) == 0).all()
         assert (np.asarray(grad(s.sum(), a).eval()) == 0).all()
 
+    def test_join_matrix1_using_column_stack(self):
+        av = np.array([0.1, 0.2, 0.3], dtype="float32")
+        bv = np.array([0.7, 0.8, 0.9], dtype="float32")
+        a = self.shared(av)
+        b = as_tensor_variable(bv)
+        s = column_stack((a, b))
+        want = np.array(
+            [[0.1, 0.7], [0.2, 0.8], [0.3, 0.9]],
+            dtype="float32",
+        )
+        out = self.eval_outputs_and_check_join([s])
+        np.testing.assert_array_equal(out, want, strict=True)
+
     def test_join_matrix1_using_vertical_stack(self):
         a = self.shared(np.array([[1, 2, 3], [4, 5, 6]], dtype=self.floatX))
         b = as_tensor_variable(np.array([[7, 8, 9]], dtype=self.floatX))
         c = as_tensor_variable(np.array([[9, 8, 7]], dtype=self.floatX))
-        s = vertical_stack(a, b, c)
+        with pytest.warns(FutureWarning):
+            s = vertical_stack(a, b, c)
 
         want = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [9, 8, 7]])
         out = self.eval_outputs_and_check_join([s])
@@ -1684,7 +1702,8 @@ def test_join_matrix1_using_horizontal_stack(self):
         a = self.shared(av)
         b = as_tensor_variable(bv)
         c = as_tensor_variable(cv)
-        s = horizontal_stack(a, b, c)
+        with pytest.warns(FutureWarning):
+            s = horizontal_stack(a, b, c)
         want = np.array(
             [[0.1, 0.2, 0.3, 0.7, 0.3, 0.2, 0.1], [0.4, 0.5, 0.6, 0.8, 0.6, 0.5, 0.4]],
             dtype="float32",
@@ -4489,15 +4508,35 @@ def test_full_like(inp, shape):
     )
 
 
+@pytest.mark.parametrize("func", [hstack, vstack, dstack])
+@pytest.mark.parametrize("dimension", [1, 2, 3])
+def test_stack_helpers(func, dimension):
+    if dimension == 1:
+        arrays = [np.arange(i * dimension, (i + 1) * dimension) for i in range(3)]
+    else:
+        arrays = [
+            np.arange(
+                i * dimension * dimension, (i + 1) * dimension * dimension
+            ).reshape(dimension, dimension)
+            for i in range(3)
+        ]
+
+    result = func(arrays)
+    np_result = getattr(np, func.__name__)(arrays)
+
+    assert np.array_equal(result.eval(), np_result)
+
+
 @pytest.mark.parametrize("func", [horizontal_stack, vertical_stack])
 def test_oriented_stack_functions(func):
-    with pytest.raises(ValueError):
-        func()
+    with pytest.warns(FutureWarning):
+        with pytest.raises(ValueError):
+            func()
 
-    a = ptb.tensor(dtype=np.float64, shape=(None, None, None))
+        a = ptb.tensor(dtype=np.float64, shape=(None, None, None))
 
-    with pytest.raises(ValueError):
-        func(a, a)
+        with pytest.raises(ValueError):
+            func(a, a)
 
 
 def test_trace():