Move contents of pennylane.utils to appropriate modules (#6588)

**Context:**
The legacy remains `qml.utils` module has been hanging around for long.
Now, finally we are to delete them and move the remaining folks to
wherever they're supposed to be.
 Specifically, the
following 4 sets of functions have been either moved or removed:

* `qml.utils._flatten`, `qml.utils.unflatten` has been moved and renamed
to
   `qml.pytrees.flatten_np` and `qml.pytrees.unflatten_np` respectively.

  * `qml.utils._inv_dict` and `qml._get_default_args` have been removed.

  * `qml.utils.pauli_eigs` has been moved to `qml.pauli.utils`.

* `qml.utils.expand_vector` has been moved to `qml.math.expand_vector`.
**Description of the Change:**

**Benefits:**
Less redundancy

**Possible Drawbacks:**
Rare chance that some downstreaming repo might be using these
funcationalities:

 - [x] lightning
 - [x] catalyst
 - [x] qml
 - [x] plugins

We will come back to check them one by one to make sure nothing is to
break

**Related GitHub Issues:**

**Related Shortcut Stories:**
[sc-76906]

---------

Co-authored-by: lillian542 <[email protected]>

doc/index.rst

@@ -230,6 +230,5 @@ PennyLane is **free** and **open source**, released under the Apache License, Ve
    code/qml_operation
    code/qml_queuing
    code/qml_tape
-   code/qml_utils
    code/qml_wires
    code/qml_workflow

doc/releases/changelog-dev.md

@@ -66,6 +66,17 @@
 
 <h3>Breaking changes 💔</h3>
 
+* The developer-facing `qml.utils` module has been removed. Specifically, the
+following 4 sets of functions have been either moved or removed[(#6588)](https://github.com/PennyLaneAI/pennylane/pull/6588):
+
+  * `qml.utils._flatten`, `qml.utils.unflatten` has been moved and renamed to `qml.optimize.qng._flatten_np` and `qml.optimize.qng._unflatten_np` respectively.
+
+  * `qml.utils._inv_dict` and `qml._get_default_args` have been removed.
+
+  * `qml.utils.pauli_eigs` has been moved to `qml.pauli.utils`.
+
+  * `qml.utils.expand_vector` has been moved to `qml.math.expand_vector`.
+
 * The `qml.qinfo` module has been removed. Please see the respective functions in the `qml.math` and `qml.measurements`
   modules instead.
   [(#6584)](https://github.com/PennyLaneAI/pennylane/pull/6584)

pennylane/math/__init__.py

@@ -34,7 +34,7 @@
 import autoray as ar
 
 from .is_independent import is_independent
-from .matrix_manipulation import expand_matrix, reduce_matrices, get_batch_size
+from .matrix_manipulation import expand_matrix, expand_vector, reduce_matrices, get_batch_size
 from .multi_dispatch import (
     add,
     array,
@@ -152,6 +152,7 @@ def __getattr__(name):
     "dot",
     "einsum",
     "expand_matrix",
+    "expand_vector",
     "expectation_value",
     "eye",
     "fidelity",

pennylane/math/matrix_manipulation.py

@@ -14,6 +14,7 @@
 """This module contains methods to expand the matrix representation of an operator
 to a higher hilbert space with re-ordered wires."""
 import itertools
+import numbers
 from collections.abc import Callable, Generator, Iterable
 from functools import reduce
 
@@ -348,3 +349,56 @@ def get_batch_size(tensor, expected_shape, expected_size):
             raise err
 
     return None
+
+
+def expand_vector(vector, original_wires, expanded_wires):
+    r"""Expand a vector to more wires.
+
+    Args:
+        vector (array): :math:`2^n` vector where n = len(original_wires).
+        original_wires (Sequence[int]): original wires of vector
+        expanded_wires (Union[Sequence[int], int]): expanded wires of vector, can be shuffled
+            If a single int m is given, corresponds to list(range(m))
+
+    Returns:
+        array: :math:`2^m` vector where m = len(expanded_wires).
+    """
+    if len(original_wires) == 0:
+        val = qml.math.squeeze(vector)
+        return val * qml.math.ones(2 ** len(expanded_wires))
+    if isinstance(expanded_wires, numbers.Integral):
+        expanded_wires = list(range(expanded_wires))
+
+    N = len(original_wires)
+    M = len(expanded_wires)
+    D = M - N
+
+    len_vector = qml.math.shape(vector)[0]
+    qudit_order = int(2 ** (np.log2(len_vector) / N))
+
+    if not set(expanded_wires).issuperset(original_wires):
+        raise ValueError("Invalid target subsystems provided in 'original_wires' argument.")
+
+    if qml.math.shape(vector) != (qudit_order**N,):
+        raise ValueError(f"Vector parameter must be of length {qudit_order}**len(original_wires)")
+
+    dims = [qudit_order] * N
+    tensor = qml.math.reshape(vector, dims)
+
+    if D > 0:
+        extra_dims = [qudit_order] * D
+        ones = qml.math.ones(qudit_order**D).reshape(extra_dims)
+        expanded_tensor = qml.math.tensordot(tensor, ones, axes=0)
+    else:
+        expanded_tensor = tensor
+
+    wire_indices = [expanded_wires.index(wire) for wire in original_wires]
+    wire_indices = np.array(wire_indices)
+
+    # Order tensor factors according to wires
+    original_indices = np.array(range(N))
+    expanded_tensor = qml.math.moveaxis(
+        expanded_tensor, tuple(original_indices), tuple(wire_indices)
+    )
+
+    return qml.math.reshape(expanded_tensor, (qudit_order**M,))

pennylane/operation.py

@@ -264,7 +264,6 @@
 from pennylane.wires import Wires, WiresLike
 
 from .pytrees import register_pytree
-from .utils import pauli_eigs
 
 # =============================================================================
 # Errors
@@ -2346,7 +2345,7 @@ def eigvals(self):
         standard_observables = {"PauliX", "PauliY", "PauliZ", "Hadamard"}
 
         # observable should be Z^{\otimes n}
-        self._eigvals_cache = pauli_eigs(len(self.wires))
+        self._eigvals_cache = qml.pauli.pauli_eigs(len(self.wires))
 
         # check if there are any non-standard observables (such as Identity)
         if set(self.name) - standard_observables:
@@ -2357,7 +2356,7 @@ def eigvals(self):
                 if k:
                     # Subgroup g contains only standard observables.
                     self._eigvals_cache = qml.math.kron(
-                        self._eigvals_cache, pauli_eigs(len(list(g)))
+                        self._eigvals_cache, qml.pauli.pauli_eigs(len(list(g)))
                     )
                 else:
                     # Subgroup g contains only non-standard observables.

pennylane/ops/op_math/composite.py

@@ -193,12 +193,12 @@ def eigvals(self):
         for ops in self.overlapping_ops:
             if len(ops) == 1:
                 eigvals.append(
-                    qml.utils.expand_vector(ops[0].eigvals(), list(ops[0].wires), list(self.wires))
+                    math.expand_vector(ops[0].eigvals(), list(ops[0].wires), list(self.wires))
                 )
             else:
                 tmp_composite = self.__class__(*ops)
                 eigvals.append(
-                    qml.utils.expand_vector(
+                    math.expand_vector(
                         tmp_composite.eigendecomposition["eigval"],
                         list(tmp_composite.wires),
                         list(self.wires),

pennylane/ops/op_math/linear_combination.py

@@ -471,12 +471,12 @@ def eigvals(self):
         for ops in self.overlapping_ops:
             if len(ops) == 1:
                 eigvals.append(
-                    qml.utils.expand_vector(ops[0].eigvals(), list(ops[0].wires), list(self.wires))
+                    qml.math.expand_vector(ops[0].eigvals(), list(ops[0].wires), list(self.wires))
                 )
             else:
                 tmp_composite = Sum(*ops)  # only change compared to CompositeOp.eigvals()
                 eigvals.append(
-                    qml.utils.expand_vector(
+                    qml.math.expand_vector(
                         tmp_composite.eigendecomposition["eigval"],
                         list(tmp_composite.wires),
                         list(self.wires),

pennylane/ops/qubit/non_parametric_ops.py

@@ -27,7 +27,6 @@
 import pennylane as qml
 from pennylane.operation import Observable, Operation
 from pennylane.typing import TensorLike
-from pennylane.utils import pauli_eigs
 from pennylane.wires import Wires, WiresLike
 
 INV_SQRT2 = 1 / qml.math.sqrt(2)
@@ -126,7 +125,7 @@ def compute_eigvals() -> np.ndarray:  # pylint: disable=arguments-differ
         >>> print(qml.Hadamard.compute_eigvals())
         [ 1 -1]
         """
-        return pauli_eigs(1)
+        return qml.pauli.pauli_eigs(1)
 
     @staticmethod
     def compute_diagonalizing_gates(wires: WiresLike) -> list[qml.operation.Operator]:
@@ -315,7 +314,7 @@ def compute_eigvals() -> np.ndarray:  # pylint: disable=arguments-differ
         >>> print(qml.X.compute_eigvals())
         [ 1 -1]
         """
-        return pauli_eigs(1)
+        return qml.pauli.pauli_eigs(1)
 
     @staticmethod
     def compute_diagonalizing_gates(wires: WiresLike) -> list[qml.operation.Operator]:
@@ -506,7 +505,7 @@ def compute_eigvals() -> np.ndarray:  # pylint: disable=arguments-differ
         >>> print(qml.Y.compute_eigvals())
         [ 1 -1]
         """
-        return pauli_eigs(1)
+        return qml.pauli.pauli_eigs(1)
 
     @staticmethod
     def compute_diagonalizing_gates(wires: WiresLike) -> list[qml.operation.Operator]:
@@ -695,7 +694,7 @@ def compute_eigvals() -> np.ndarray:  # pylint: disable=arguments-differ
         >>> print(qml.Z.compute_eigvals())
         [ 1 -1]
         """
-        return pauli_eigs(1)
+        return qml.pauli.pauli_eigs(1)
 
     @staticmethod
     def compute_diagonalizing_gates(  # pylint: disable=unused-argument

pennylane/ops/qubit/parametric_ops_multi_qubit.py

@@ -27,7 +27,6 @@
 from pennylane.math import expand_matrix
 from pennylane.operation import AnyWires, FlatPytree, Operation
 from pennylane.typing import TensorLike
-from pennylane.utils import pauli_eigs
 from pennylane.wires import Wires, WiresLike
 
 from .non_parametric_ops import Hadamard, PauliX, PauliY, PauliZ
@@ -105,7 +104,7 @@ def compute_matrix(
                 [0.0000+0.0000j, 0.0000+0.0000j, 0.9988+0.0500j, 0.0000+0.0000j],
                 [0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j, 0.9988-0.0500j]])
         """
-        eigs = qml.math.convert_like(pauli_eigs(num_wires), theta)
+        eigs = qml.math.convert_like(qml.pauli.pauli_eigs(num_wires), theta)
 
         if qml.math.get_interface(theta) == "tensorflow":
             theta = qml.math.cast_like(theta, 1j)
@@ -153,7 +152,7 @@ def compute_eigvals(
         tensor([0.9689-0.2474j, 0.9689+0.2474j, 0.9689+0.2474j, 0.9689-0.2474j,
                 0.9689+0.2474j, 0.9689-0.2474j, 0.9689-0.2474j, 0.9689+0.2474j])
         """
-        eigs = qml.math.convert_like(pauli_eigs(num_wires), theta)
+        eigs = qml.math.convert_like(qml.pauli.pauli_eigs(num_wires), theta)
 
         if qml.math.get_interface(theta) == "tensorflow":
             theta = qml.math.cast_like(theta, 1j)

pennylane/optimize/momentum_qng.py

@@ -16,9 +16,8 @@
 # pylint: disable=too-many-branches
 # pylint: disable=too-many-arguments
 from pennylane import numpy as pnp
-from pennylane.utils import _flatten, unflatten
 
-from .qng import QNGOptimizer
+from .qng import QNGOptimizer, _flatten_np, _unflatten_np
 
 
 class MomentumQNGOptimizer(QNGOptimizer):
@@ -131,12 +130,12 @@ def apply_grad(self, grad, args):
 
         for index, arg in enumerate(args):
             if getattr(arg, "requires_grad", False):
-                grad_flat = pnp.array(list(_flatten(grad[trained_index])))
+                grad_flat = pnp.array(list(_flatten_np(grad[trained_index])))
                 # self.metric_tensor has already been reshaped to 2D, matching flat gradient.
                 qng_update = pnp.linalg.pinv(metric_tensor[trained_index]) @ grad_flat
 
                 self.accumulation[trained_index] *= self.momentum
-                self.accumulation[trained_index] += self.stepsize * unflatten(
+                self.accumulation[trained_index] += self.stepsize * _unflatten_np(
                     qng_update, grad[trained_index]
                 )
                 args_new[index] = arg - self.accumulation[trained_index]

pennylane/optimize/qng.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Quantum natural gradient optimizer"""
+import numbers
+from collections.abc import Iterable
+
 import pennylane as qml
 
 # pylint: disable=too-many-branches
 # pylint: disable=too-many-arguments
 from pennylane import numpy as pnp
-from pennylane.utils import _flatten, unflatten
 
 from .gradient_descent import GradientDescentOptimizer
 
@@ -277,11 +279,88 @@ def apply_grad(self, grad, args):
         trained_index = 0
         for index, arg in enumerate(args):
             if getattr(arg, "requires_grad", False):
-                grad_flat = pnp.array(list(_flatten(grad[trained_index])))
+                grad_flat = pnp.array(list(_flatten_np(grad[trained_index])))
                 # self.metric_tensor has already been reshaped to 2D, matching flat gradient.
                 update = pnp.linalg.pinv(mt[trained_index]) @ grad_flat
-                args_new[index] = arg - self.stepsize * unflatten(update, grad[trained_index])
+                args_new[index] = arg - self.stepsize * _unflatten_np(update, grad[trained_index])
 
                 trained_index += 1
 
         return tuple(args_new)
+
+
+def _flatten_np(x):
+    """Iterate recursively through an arbitrarily nested structure in depth-first order.
+
+    See also :func:`_unflatten`.
+
+    Args:
+        x (array, Iterable, Any): each element of an array or an Iterable may itself be any of these types
+
+    Yields:
+        Any: elements of x in depth-first order
+    """
+    if isinstance(x, pnp.ndarray):
+        yield from _flatten_np(
+            x.flat
+        )  # should we allow object arrays? or just "yield from x.flat"?
+    elif isinstance(x, qml.wires.Wires):
+        # Reursive calls to flatten `Wires` will cause infinite recursion (`Wires` atoms are `Wires`).
+        # Since Wires are always flat, just yield.
+        yield from x
+    elif isinstance(x, Iterable) and not isinstance(x, (str, bytes)):
+        for item in x:
+            yield from _flatten_np(item)
+    else:
+        yield x
+
+
+def _unflatten_np_dispatch(flat, model):
+    """Restores an arbitrary nested structure to a flattened iterable.
+
+    See also :func:`_flatten`.
+
+    Args:
+        flat (array): 1D array of items
+        model (array, Iterable, Number): model nested structure
+
+    Raises:
+        TypeError: if ``model`` contains an object of unsupported type
+
+    Returns:
+        Union[array, list, Any], array: first elements of flat arranged into the nested
+        structure of model, unused elements of flat
+    """
+    if isinstance(model, (numbers.Number, str)):
+        return flat[0], flat[1:]
+
+    if isinstance(model, pnp.ndarray):
+        idx = model.size
+        res = pnp.array(flat)[:idx].reshape(model.shape)
+        return res, flat[idx:]
+
+    if isinstance(model, Iterable):
+        res = []
+        for x in model:
+            val, flat = _unflatten_np_dispatch(flat, x)
+            res.append(val)
+        return res, flat
+
+    raise TypeError(f"Unsupported type in the model: {type(model)}")
+
+
+def _unflatten_np(flat, model):
+    """Wrapper for :func:`_unflatten`.
+
+    Args:
+        flat (array): 1D array of items
+        model (array, Iterable, Number): model nested structure
+
+    Raises:
+        ValueError: if ``flat`` has more elements than ``model``
+    """
+    # pylint:disable=len-as-condition
+    res, tail = _unflatten_np_dispatch(pnp.asarray(flat), model)
+    if len(tail) != 0:
+        raise ValueError("Flattened iterable has more elements than the model.")
+    return res