Simplify pack/unpack calls

Author: ricardoV94

pytensor/tensor/optimize.py

@@ -447,8 +447,8 @@ def compute_implicit_gradients(
             )
             grad_wrt_args = unpack(
                 grad_wrt_args_packed,
-                packed_shapes=packed_shapes_from_outer,
-                axes=0 if not all(inp.ndim == 0 for inp in (x_star, *args)) else None,
+                packed_shapes_from_outer,
+                keep_axes=None if all(inp.ndim == 0 for inp in (x_star, *args)) else 0,
             )
         else:
             grad_wrt_args = [grad_wrt_args_packed]
@@ -702,8 +702,8 @@ def pack_inputs_of_objective(
     elif len(x) == 1:
         packed_input = x[0]
     else:
-        packed_input, packed_shapes = pack(*x, axes=None)
-        unpacked_output = unpack(packed_input, axes=None, packed_shapes=packed_shapes)
+        packed_input, packed_shapes = pack(*x)
+        unpacked_output = unpack(packed_input, packed_shapes)
 
         objective = graph_replace(
             objective,
@@ -781,7 +781,7 @@ def minimize(
     solution, success = minimize_op(packed_input, *args)
 
     if packed_shapes is not None:
-        solution = unpack(solution, axes=None, packed_shapes=packed_shapes)
+        solution = unpack(solution, packed_shapes)
 
     return solution, success
 
@@ -1109,7 +1109,7 @@ def root(
 
     solution, success = root_op(packed_variables, *args)
     if packed_shapes is not None:
-        solution = unpack(solution, axes=None, packed_shapes=packed_shapes)
+        solution = unpack(solution, packed_shapes)
 
     return solution, success