.wip

Author: jessegrabowski

pytensor/tensor/extra_ops.py

@@ -2259,7 +2259,9 @@ class Pack(OpFromGraph):
 
 
 def pack(
-    *tensors: TensorVariable, axes: int | Sequence[int] | None = None
+    *tensors: TensorVariable,
+    axes: int | Sequence[int] | None = None,
+    preserve_axes: int | Sequence[int] | None = None,
 ) -> tuple[TensorVariable, list[tuple[TensorVariable]]]:
     """
     Given a list of tensors of varying shapes and dimensions, ravels and concatenates them into a single 1d vector.
@@ -2270,10 +2272,11 @@ def pack(
         Tensors to be packed. Tensors can have varying shapes and dimensions, but must have the same size along each
         of the dimensions specified in the `axes` parameter.
     axes: int or sequence of int, optional
-        Axes to be preserved. All other axes will be raveled (packed), and the output is the result of concatenating
-        on the new raveled dimension. If None, all axes will be raveled and joined. Axes can be either positive or
+        Axes to be raveled. All other axes will preserved, and the output is the result of concatenating
+        on the raveled dimension. If None, all axes will be raveled and joined. Axes can be either positive or
         negative, but must be striclty increasing in both the positive and  negative parts of the list. Negative axes
-        must come after positive axes.
+        must come after positive axes. Only one of `axes` or `preserve_axes` can be specified.
+    preserve_axes: int or sequence of int, optional
 
     Returns
     -------