Add unpack helper

Author: jessegrabowski

pytensor/tensor/extra_ops.py

@@ -26,7 +26,7 @@
 from pytensor.scalar import upcast
 from pytensor.tensor import TensorLike, as_tensor_variable
 from pytensor.tensor import basic as ptb
-from pytensor.tensor.basic import alloc, as_tensor, join, second
+from pytensor.tensor.basic import alloc, as_tensor, join, second, split
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import abs as pt_abs
 from pytensor.tensor.math import all as pt_all
@@ -2198,6 +2198,35 @@ def pack(*tensors: TensorLike, axes: Sequence[int] | int | None = None):
     return join(n_before, *reshaped_tensors), packed_shapes
 
 
+def unpack(packed_input, axes, packed_shapes):
+    if axes is None:
+        if packed_input.ndim != 1:
+            raise ValueError(
+                "unpack can only be called with keep_axis=None for 1d inputs"
+            )
+        split_axis = 0
+    else:
+        axes = normalize_axis_tuple(axes, ndim=packed_input.ndim)
+        try:
+            [split_axis] = (i for i in range(packed_input.ndim) if i not in axes)
+        except ValueError as err:
+            raise ValueError(
+                "Unpack must have exactly one more dimension that implied by axes"
+            ) from err
+
+    split_inputs = split(
+        packed_input,
+        splits_size=[prod(shape).astype(int) for shape in packed_shapes],
+        n_splits=len(packed_shapes),
+        axis=split_axis,
+    )
+
+    return [
+        split_dims(inp, shape, split_axis)
+        for inp, shape in zip(split_inputs, packed_shapes, strict=True)
+    ]
+
+
 __all__ = [
     "bartlett",
     "bincount",

tests/tensor/test_extra_ops.py

@@ -47,6 +47,7 @@
     split_dims,
     squeeze,
     to_one_hot,
+    unpack,
     unravel_index,
 )
 from pytensor.tensor.type import (
@@ -1593,3 +1594,50 @@ def test_pack_basic(self):
                 packed_shape.eval(input_dict, on_unused_input="ignore"),
                 tensor.type.shape[1:-1],
             )
+
+    @pytest.mark.parametrize("axes", [None, -1, (-2, -1)])
+    def test_pack_unpack_round_trip(self, axes):
+        rng = np.random.default_rng()
+
+        x = pt.tensor("x", shape=(3, 5))
+        y = pt.tensor("y", shape=(3, 3, 5))
+        z = pt.tensor("z", shape=(1, 3, 5))
+
+        flat_packed, packed_shapes = pack(x, y, z, axes=axes)
+        new_outputs = unpack(flat_packed, axes=axes, packed_shapes=packed_shapes)
+
+        fn = pytensor.function([x, y, z], new_outputs, mode="FAST_COMPILE")
+
+        input_dict = {
+            var.name: rng.normal(size=var.type.shape).astype(config.floatX)
+            for var in [x, y, z]
+        }
+        output_vals = fn(**input_dict)
+
+        for input_val, output_val in zip(input_dict.values(), output_vals, strict=True):
+            np.testing.assert_allclose(input_val, output_val)
+
+
+def test_make_replacements_with_pack_unpack():
+    rng = np.random.default_rng()
+
+    x = pt.tensor("x", shape=())
+    y = pt.tensor("y", shape=(5,))
+    z = pt.tensor("z", shape=(3, 3))
+
+    loss = (x + y.sum() + z.sum()) ** 2
+
+    flat_packed, packed_shapes = pack(x, y, z, axes=None)
+    new_input = flat_packed.type()
+    new_outputs = unpack(new_input, axes=None, packed_shapes=packed_shapes)
+
+    loss = pytensor.graph.graph_replace(loss, dict(zip([x, y, z], new_outputs)))
+    fn = pytensor.function([new_input, x, y, z], loss, mode="FAST_COMPILE")
+
+    input_vals = [
+        rng.normal(size=(var.type.shape)).astype(config.floatX) for var in [x, y, z]
+    ]
+    flat_inputs = np.concatenate([input.ravel() for input in input_vals], axis=0)
+    output_val = fn(flat_inputs, *input_vals)
+
+    assert np.allclose(output_val, sum([input.sum() for input in input_vals]) ** 2)