Repository: PennyLaneAI/pennylane
Pull request #2168 (2022-02-04 03:43:02)

Nailing down automatic circuit cutter's API
**Context:**
When passing a tape to the automatic circuit cutter (KaHyPar), several hyperparameters (e.g. `k`, `imbalance`, etc.) need to be supplied. Depending on use cases, sensible default values should be derived based on known information (which usually comes at different points in time) such as the available devices and the circuits. However, power users should also be able to override and supply their own parameters. We thus must provide a unified interface that's general enough to allow these types of interactions.

Drawing similarities between distributing circuit fragments to devices and the classical problem of distributing computation tasks dags to multiple workers, the constraints coming from devices should be separated from the requirements coming from the computations. This separation is realized inside the added class `CutStrategy`, whose attributes contain constraints from devices and whose method `get_cut_kwargs()` takes circuit requirements at runtime and returns complete sets of cutting hyperparameters. In terms of UI, both the device constrains and the calling of `get_cut_kwargs()` could be done either manually by users or within decorators. `dataclass` was elected to make the implementation of the init function cleaner.

**Description of the Change:**
Adds:
- a class `CutStrategy` that can be used as an interface to coordinate devices info, user constraints, and circuit requirements into  partition hyperparameters ready for partitioning methods.

**Possible Drawbacks:**
Gate cut compatibility not fully explored


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2168](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (7e25adc) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/5a28983fc7bd950cde8a4014e54261fef4b54293?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (5a28983) will **decrease** coverage by `3.20%`.
> The diff coverage is `31.19%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #2168      +/-   ##
==========================================
- Coverage   99.20%   96.00%   -3.21%     
==========================================
  Files         231      231              
  Lines       17785    17894     +109     
==========================================
- Hits        17644    17179     -465     
- Misses        141      715     +574     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/qcut.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvcWN1dC5weQ==) | `56.89% <31.19%> (-43.11%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax\_jit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4X2ppdC5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.63%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `77.08% <0.00%> (-22.92%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `84.41% <0.00%> (-15.59%)` | :arrow_down: |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `87.73% <0.00%> (-12.27%)` | :arrow_down: |
| [pennylane/math/is\_independent.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvaXNfaW5kZXBlbmRlbnQucHk=) | `88.50% <0.00%> (-11.50%)` | :arrow_down: |
| ... and [14 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [5a28983...7e25adc](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2168?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

trbromley: [sc-14152]
zeyueN: Edited the description. Promoting to real PR.
zeyueN: @trbromley I think I've addressed all the comments and ready for review.

The documentation check workflow is failing for some reason. Any idea how to fix?

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 9491b3e35..eddf4304e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -81,6 +81,11 @@
   The postprocessing function for the `cut_circuit` transform has been added.
   [(#2192)](https://github.com/PennyLaneAI/pennylane/pull/2192)
 
+  A class `CutStrategy` which acts as an interface and coordinates device/user
+  constraints with circuit execution requirements to come up with the best sets
+  of graph partitioning parameters.
+  [(#2168)](https://github.com/PennyLaneAI/pennylane/pull/2168)
+
 <h3>Improvements</h3>
 
 * The `gradients` module has been streamlined and special-purpose functions
@@ -288,6 +293,6 @@ The Operator class has undergone a major refactor with the following changes:
 This release contains contributions from (in alphabetical order):
 
 Thomas Bromley, Anthony Hayes, Josh Izaac, Christina Lee,
-Maria Fernanda Morris, Maria Schuld, Jay Soni, Antal Száva,
+Maria Fernanda Morris, Zeyue Niu, Maria Schuld, Jay Soni, Antal Száva,
 David Wierichs
 
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 57855710a..9b91b90de 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -98,6 +98,7 @@ The following are utility functions that compose the circuit cutting transform.
     ~transforms.expand_fragment_tapes
     ~transforms.contract_tensors
     ~transforms.qcut_processing_fn
+    ~transforms.CutStrategy
 
 Transforms that act on tapes
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -179,4 +180,5 @@ from .qcut import (
     expand_fragment_tapes,
     contract_tensors,
     qcut_processing_fn,
+    CutStrategy,
 )
diff --git a/pennylane/transforms/qcut.py b/pennylane/transforms/qcut.py
index 45b12bceb..b65514f4e 100644
--- a/pennylane/transforms/qcut.py
+++ b/pennylane/transforms/qcut.py
@@ -18,9 +18,11 @@ circuits to be distributed across multiple devices.
 
 import copy
 import string
+import warnings
 import uuid
+from typing import Sequence, Tuple, List, Dict, Any, Union, ClassVar
 from itertools import product
-from typing import List, Sequence, Tuple
+from dataclasses import dataclass, InitVar
 
 import pennylane as qml
 from networkx import MultiDiGraph, weakly_connected_components
@@ -853,3 +855,324 @@ def qcut_processing_fn(
         tensors, communication_graph, prepare_nodes, measure_nodes, use_opt_einsum
     )
     return result
+
+
+@dataclass()
+class CutStrategy:
+    """
+    A circuit-cutting distribution policy for executing (large) circuits on available (comparably
+    smaller) devices.
+
+    Args:
+        devices (Union[qml.Device, Sequence[qml.Device]]): Single, or Sequence of, device(s).
+            Optional only when ``max_free_wires`` is provided.
+        max_free_wires (int): Number of wires for the largest available device. Optional only when
+            ``devices`` is provided where it defaults to the maximum number of wires among
+            ``devices``.
+        min_free_wires (int): Number of wires for the smallest available device, or, equivalently,
+            the smallest max fragment-wire-size that the partitioning is allowed to explore.
+            When provided, this parameter will be used to derive an upper-bound to the range of
+            explored number of fragments.  Optional, defaults to ``max_free_wires``.
+        num_fragments_probed (Union[int, Sequence[int]]): Single, or 2-Sequence of, number(s)
+            specifying the potential (range of) number of fragments for the partitioner to attempt.
+            Optional, defaults to probing all valid strategies derivable from the circuit and
+            devices.
+        max_free_gates (int): Maximum allowed circuit depth for the deepest available device.
+            Optional, defaults to unlimited depth.
+        min_free_gates (int): Maximum allowed circuit depth for the shallowest available device.
+            Optional, defaults to ``max_free_gates``.
+        imbalance_tolerance (float): The global maximum allowed imbalance for all partition trials.
+            Optional, defaults to unlimited imbalance. Used only if there's a known hard balancing
+            constraint on the partitioning problem.
+
+    **Example**
+
+    The following cut strategy specifies that a circuit should be cut into between
+    ``2`` to ``5`` fragments, with each fragment having at most ``6`` wires and
+    at least ``4`` wires:
+
+    >>> cut_strategy = qml.transforms.CutStrategy(
+    ...     max_free_wires=6,
+    ...     min_free_wires=4,
+    ...     num_fragments_probed=(2, 5),
+    ... )
+
+    """
+
+    # pylint: disable=too-many-arguments, too-many-instance-attributes
+
+    #: Initialization argument only, used to derive ``max_free_wires`` and ``min_free_wires``.
+    devices: InitVar[Union[qml.Device, Sequence[qml.Device]]] = None
+
+    #: Number of wires for the largest available device.
+    max_free_wires: int = None
+    #: Number of wires for the smallest available device.
+    min_free_wires: int = None
+    #: The potential (range of) number of fragments for the partitioner to attempt.
+    num_fragments_probed: Union[int, Sequence[int]] = None
+    #: Maximum allowed circuit depth for the deepest available device.
+    max_free_gates: int = None
+    #: Maximum allowed circuit depth for the shallowest available device.
+    min_free_gates: int = None
+    #: The global maximum allowed imbalance for all partition trials.
+    imbalance_tolerance: float = None
+
+    #: Class attribute, threshold for warning about too many fragments.
+    HIGH_NUM_FRAGMENTS: ClassVar[int] = 20
+    #: Class attribute, threshold for warning about too many partition attempts.
+    HIGH_PARTITION_ATTEMPTS: ClassVar[int] = 20
+
+    def __post_init__(
+        self,
+        devices,
+    ):
+        """Deriving cutting constraints from given devices and parameters."""
+
+        self.max_free_wires = self.max_free_wires or self.min_free_wires
+        if isinstance(self.num_fragments_probed, int):
+            self.num_fragments_probed = [self.num_fragments_probed]
+        if isinstance(self.num_fragments_probed, (list, tuple)):
+            self.num_fragments_probed = sorted(self.num_fragments_probed)
+            self.k_lower = self.num_fragments_probed[0]
+            self.k_upper = self.num_fragments_probed[-1]
+            if self.k_lower <= 0:
+                raise ValueError("`num_fragments_probed` must be positive int(s)")
+        else:
+            self.k_lower, self.k_upper = None, None
+
+        if devices is None and self.max_free_wires is None:
+            raise ValueError("One of arguments `devices` and max_free_wires` must be provided.")
+
+        if isinstance(devices, qml.Device):
+            devices = (devices,)
+
+        if devices is not None:
+            if not isinstance(devices, Sequence) or any(
+                (not isinstance(d, qml.Device) for d in devices)
+            ):
+                raise ValueError(
+                    "Argument `devices` must be a list or tuple containing elements of type "
+                    "`qml.Device`"
+                )
+
+            device_wire_sizes = [len(d.wires) for d in devices]
+
+            self.max_free_wires = self.max_free_wires or max(device_wire_sizes)
+            self.min_free_wires = self.min_free_wires or min(device_wire_sizes)
+
+        if (self.imbalance_tolerance is not None) and not (
+            isinstance(self.imbalance_tolerance, (float, int)) and self.imbalance_tolerance >= 0
+        ):
+            raise ValueError(
+                "The overall `imbalance_tolerance` is expected to be a non-negative number, "
+                f"got {type(self.imbalance_tolerance)} with value {self.imbalance_tolerance}."
+            )
+
+    def get_cut_kwargs(
+        self,
+        tape_dag: MultiDiGraph,
+        max_wires_by_fragment: Sequence[int] = None,
+        max_gates_by_fragment: Sequence[int] = None,
+    ) -> List[Dict[str, Any]]:
+        """Derive the complete set of arguments, based on a given circuit, for passing to a graph
+        partitioner.
+
+        Args:
+            tape_dag (MultiDiGraph): Graph representing a tape, typically the output of
+                :func:`tape_to_graph`.
+            max_wires_by_fragment (Sequence[int]): User-predetermined list of wire limits by
+                fragment. If supplied, the number of fragments will be derived from it and
+                exploration of other choices will not be made.
+            max_gates_by_fragment (Sequence[int]): User-predetermined list of gate limits by
+                fragment. If supplied, the number of fragments will be derived from it and
+                exploration of other choices will not be made.
+
+        Returns:
+            List[Dict[str, Any]]: A list of minimal kwargs being passed to a graph
+            partitioner method.
+
+        **Example**
+
+        Deriving kwargs for a given circuit and feeding them to a custom partitioner, along with
+        extra parameters specified using ``extra_kwargs``:
+
+        >>> cut_strategy = qcut.CutStrategy(devices=dev)
+        >>> cut_kwargs = cut_strategy.get_cut_kwargs(tape_dag)
+        >>> cut_trials = [
+        ...     my_partition_fn(tape_dag, **kwargs, **extra_kwargs) for kwargs in cut_kwargs
+        ... ]
+
+        """
+        tape_wires = set(w for _, _, w in tape_dag.edges.data("wire"))
+        num_tape_wires = len(tape_wires)
+        num_tape_gates = tape_dag.order()
+        self._validate_input(max_wires_by_fragment, max_gates_by_fragment)
+
+        probed_cuts = self._infer_probed_cuts(
+            num_tape_wires=num_tape_wires,
+            num_tape_gates=num_tape_gates,
+            max_wires_by_fragment=max_wires_by_fragment,
+            max_gates_by_fragment=max_gates_by_fragment,
+        )
+
+        return probed_cuts
+
+    @staticmethod
+    def _infer_imbalance(
+        k, num_wires, num_gates, free_wires, free_gates, imbalance_tolerance=None
+    ) -> float:
+        """Helper function for determining best imbalance limit."""
+        avg_fragment_wires = (num_wires - 1) // k + 1
+        avg_fragment_gates = (num_gates - 1) // k + 1
+        if free_wires < avg_fragment_wires:
+            raise ValueError(
+                "`free_wires` should be no less than the average number of wires per fragment. "
+                f"Got {free_wires} >= {avg_fragment_wires} ."
+            )
+        if free_gates < avg_fragment_gates:
+            raise ValueError(
+                "`free_gates` should be no less than the average number of gates per fragment. "
+                f"Got {free_gates} >= {avg_fragment_gates} ."
+            )
+
+        wire_imbalance = free_wires / avg_fragment_wires - 1
+        gate_imbalance = free_gates / avg_fragment_gates - 1
+        imbalance = min(gate_imbalance, wire_imbalance)
+        if imbalance_tolerance is not None:
+            imbalance = min(imbalance, imbalance_tolerance)
+
+        return imbalance
+
+    @staticmethod
+    def _validate_input(
+        max_wires_by_fragment,
+        max_gates_by_fragment,
+    ):
+        """Helper parameter checker."""
+        if max_wires_by_fragment is not None:
+            if not isinstance(max_wires_by_fragment, (list, tuple)):
+                raise ValueError(
+                    "`max_wires_by_fragment` is expected to be a list or tuple, but got "
+                    f"{type(max_gates_by_fragment)}."
+                )
+            if any(not (isinstance(i, int) and i > 0) for i in max_wires_by_fragment):
+                raise ValueError(
+                    "`max_wires_by_fragment` is expected to contain positive integers only."
+                )
+        if max_gates_by_fragment is not None:
+            if not isinstance(max_gates_by_fragment, (list, tuple)):
+                raise ValueError(
+                    "`max_gates_by_fragment` is expected to be a list or tuple, but got "
+                    f"{type(max_gates_by_fragment)}."
+                )
+            if any(not (isinstance(i, int) and i > 0) for i in max_gates_by_fragment):
+                raise ValueError(
+                    "`max_gates_by_fragment` is expected to contain positive integers only."
+                )
+        if max_wires_by_fragment is not None and max_gates_by_fragment is not None:
+            if len(max_wires_by_fragment) != len(max_gates_by_fragment):
+                raise ValueError(
+                    "The lengths of `max_wires_by_fragment` and `max_gates_by_fragment` should be "
+                    f"equal, but got {len(max_wires_by_fragment)} and {len(max_gates_by_fragment)}."
+                )
+
+    def _infer_probed_cuts(
+        self,
+        num_tape_wires,
+        num_tape_gates,
+        max_wires_by_fragment=None,
+        max_gates_by_fragment=None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Helper function for deriving the minimal set of best default partitioning constraints
+        for the graph partitioner.
+
+        Args:
+            num_tape_wires (int): Number of wires in the circuit tape to be partitioned.
+            num_tape_gates (int): Number of gates in the circuit tape to be partitioned.
+            max_wires_by_fragment (Sequence[int]): User-predetermined list of wire limits by
+                fragment. If supplied, the number of fragments will be derived from it and
+                exploration of other choices will not be made.
+            max_gates_by_fragment (Sequence[int]): User-predetermined list of gate limits by
+                fragment. If supplied, the number of fragments will be derived from it and
+                exploration of other choices will not be made.
+
+        Returns:
+            List[Dict[str, Any]]: A list of minimal set of kwargs being passed to a graph
+                partitioner method.
+        """
+
+        # Assumes unlimited width/depth if not supplied.
+        max_free_wires = self.max_free_wires or num_tape_wires
+        max_free_gates = self.max_free_gates or num_tape_gates
+
+        # Assumes same number of wires/gates across all devices if min_free_* not provided.
+        min_free_wires = self.min_free_wires or max_free_wires
+        min_free_gates = self.min_free_gates or max_free_gates
+
+        # The lower bound of k corresponds to executing each fragment on the largest available device.
+        k_lb = 1 + max(
+            (num_tape_wires - 1) // max_free_wires,  # wire limited
+            (num_tape_gates - 1) // max_free_gates,  # gate limited
+        )
+        # The upper bound of k corresponds to executing each fragment on the smallest available device.
+        k_ub = 1 + max(
+            (num_tape_wires - 1) // min_free_wires,  # wire limited
+            (num_tape_gates - 1) // min_free_gates,  # gate limited
+        )
+
+        # The global imbalance tolerance, if not given, defaults to a very loose upper bound:
+        imbalance_tolerance = k_ub if self.imbalance_tolerance is None else self.imbalance_tolerance
+
+        probed_cuts = []
+
+        if max_gates_by_fragment is None and max_wires_by_fragment is None:
+
+            # k_lower, when supplied by a user, can be higher than k_lb if the the desired k is known:
+            k_lower = self.k_lower if self.k_lower is not None else k_lb
+            # k_upper, when supplied by a user, can be higher than k_ub to encourage exploration:
+            k_upper = self.k_upper if self.k_upper is not None else k_ub
+
+            if k_lower < k_lb:
+                warnings.warn(
+                    f"The provided `k_lower={k_lower}` is less than the lowest allowed value, "
+                    f"will override and set `k_lower={k_lb}`."
+                )
+                k_lower = k_lb
+
+            if k_lower > self.HIGH_NUM_FRAGMENTS:
+                warnings.warn(
+                    f"The attempted number of fragments seems high with lower bound at {k_lower}."
+                )
+
+            # Prepare the list of ks to explore:
+            ks = list(range(k_lower, k_upper + 1))
+
+            if len(ks) > self.HIGH_PARTITION_ATTEMPTS:
+                warnings.warn(f"The numer of partition attempts seems high ({len(ks)}).")
+        else:
+            # When the by-fragment wire and/or gate limits are supplied, derive k and imbalance and
+            # return a single partition config.
+            ks = [len(max_wires_by_fragment or max_gates_by_fragment)]
+
+        for k in ks:
+            imbalance = self._infer_imbalance(
+                k,
+                num_tape_wires,
+                num_tape_gates,
+                max_free_wires if max_wires_by_fragment is None else max(max_wires_by_fragment),
+                max_free_gates if max_gates_by_fragment is None else max(max_gates_by_fragment),
+                imbalance_tolerance,
+            )
+            cut_kwargs = {
+                "num_fragments": k,
+                "imbalance": imbalance,
+            }
+            if max_wires_by_fragment is not None:
+                cut_kwargs["max_wires_by_fragment"] = max_wires_by_fragment
+            if max_gates_by_fragment is not None:
+                cut_kwargs["max_gates_by_fragment"] = max_gates_by_fragment
+
+            probed_cuts.append(cut_kwargs)
+
+        return probed_cuts


Test patch:
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index b12df8ae3..bfdecd53f 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -1799,3 +1799,188 @@ class TestQCutProcessingFn:
         ) * f(1)
 
         assert np.allclose(grad, expected_grad)
+
+
+class TestCutStrategy:
+    """Tests for class CutStrategy"""
+
+    devs = [qml.device("default.qubit", wires=n) for n in [4, 6]]
+    tape_dags = [qcut.tape_to_graph(t) for t in [tape, multi_cut_tape]]
+
+    @pytest.mark.parametrize("devices", [None, 1, devs[0]])
+    @pytest.mark.parametrize("imbalance_tolerance", [None, -1])
+    @pytest.mark.parametrize("num_fragments_probed", [None, 0])
+    def test_init_raises(self, devices, imbalance_tolerance, num_fragments_probed):
+        """Test if ill-initialized instances throw errors."""
+
+        if (
+            isinstance(devices, qml.Device)
+            and imbalance_tolerance is None
+            and num_fragments_probed is None
+        ):
+            return  # skip the only valid combination
+
+        with pytest.raises(ValueError):
+            qcut.CutStrategy(
+                devices=devices,
+                num_fragments_probed=num_fragments_probed,
+                imbalance_tolerance=imbalance_tolerance,
+            )
+
+    @pytest.mark.parametrize("devices", [devs[0], devs])
+    @pytest.mark.parametrize("max_free_wires", [None, 3])
+    @pytest.mark.parametrize("min_free_wires", [None, 2])
+    @pytest.mark.parametrize("num_fragments_probed", [None, 2, (2, 4)])
+    @pytest.mark.parametrize("imbalance_tolerance", [None, 0, 0.1])
+    def test_init(
+        self, devices, max_free_wires, min_free_wires, num_fragments_probed, imbalance_tolerance
+    ):
+        """Test the __post_init__ properly sets defaults based on provided info."""
+
+        strategy = qcut.CutStrategy(
+            devices=devices,
+            max_free_wires=max_free_wires,
+            num_fragments_probed=num_fragments_probed,
+            imbalance_tolerance=imbalance_tolerance,
+        )
+
+        devices = [devices] if not isinstance(devices, list) else devices
+
+        max_dev_wires = max((len(d.wires) for d in devices))
+        assert strategy.max_free_wires == max_free_wires or max_dev_wires or min_free_wires
+        assert strategy.min_free_wires == min_free_wires or max_free_wires or max_dev_wires
+        assert strategy.imbalance_tolerance == imbalance_tolerance
+
+        if num_fragments_probed is not None:
+            assert (
+                strategy.k_lower == num_fragments_probed
+                if isinstance(num_fragments_probed, int)
+                else min(num_fragments_probed)
+            )
+            assert (
+                strategy.k_upper == num_fragments_probed
+                if isinstance(num_fragments_probed, int)
+                else max(num_fragments_probed)
+            )
+        else:
+            assert strategy.k_lower is None
+            assert strategy.k_upper is None
+
+    @pytest.mark.parametrize("k", [4, 5, 6])
+    @pytest.mark.parametrize("imbalance_tolerance", [None, 0, 0.1])
+    def test_infer_wire_imbalance(self, k, imbalance_tolerance):
+        """Test that the imbalance is correctly derived under simple circumstances."""
+
+        num_wires = 10
+        num_gates = 10
+        free_wires = 3
+
+        imbalance = qcut.CutStrategy._infer_imbalance(
+            k=k,
+            num_wires=num_wires,
+            num_gates=num_gates,
+            free_wires=free_wires,
+            free_gates=1000,
+            imbalance_tolerance=imbalance_tolerance,
+        )
+
+        avg_size = int(num_wires / k + 1 - 1e-7)
+        if imbalance_tolerance is not None:
+            assert imbalance <= imbalance_tolerance
+        else:
+            assert imbalance == free_wires / avg_size - 1
+
+    @pytest.mark.parametrize("num_wires", [50, 10])
+    def test_infer_wire_imbalance_raises(
+        self,
+        num_wires,
+    ):
+        """Test that the imbalance correctly raises."""
+
+        k = 2
+        num_gates = 50
+
+        with pytest.raises(ValueError, match=f"`free_{'wires' if num_wires > 40 else 'gates'}`"):
+            qcut.CutStrategy._infer_imbalance(
+                k=k,
+                num_wires=num_wires,
+                num_gates=num_gates,
+                free_wires=20,
+                free_gates=20,
+            )
+
+    @pytest.mark.parametrize("devices", [devs[0], devs])
+    @pytest.mark.parametrize("num_fragments_probed", [None, 4, (4, 6)])
+    @pytest.mark.parametrize("imbalance_tolerance", [None, 0, 0.1])
+    @pytest.mark.parametrize("tape_dag", tape_dags)
+    def test_get_cut_kwargs(self, devices, num_fragments_probed, imbalance_tolerance, tape_dag):
+        """Test that the cut kwargs can be derived."""
+
+        strategy = qcut.CutStrategy(
+            devices=devices,
+            num_fragments_probed=num_fragments_probed,
+            imbalance_tolerance=imbalance_tolerance,
+        )
+
+        all_cut_kwargs = strategy.get_cut_kwargs(tape_dag=tape_dag)
+
+        assert all_cut_kwargs
+        assert all("imbalance" in kwargs and "num_fragments" in kwargs for kwargs in all_cut_kwargs)
+        if imbalance_tolerance is not None:
+            assert all([kwargs["imbalance"] <= imbalance_tolerance for kwargs in all_cut_kwargs])
+
+    @pytest.mark.parametrize(
+        "num_fragments_probed", [1, qcut.CutStrategy.HIGH_NUM_FRAGMENTS + 1, (2, 100)]
+    )
+    def test_get_cut_kwargs_warnings(self, num_fragments_probed):
+        """Test the 3 situations where the get_cut_kwargs pops out a warning."""
+        strategy = qcut.CutStrategy(
+            max_free_wires=2,
+            num_fragments_probed=num_fragments_probed,
+        )
+        k = num_fragments_probed
+        k_lower = k if isinstance(k, int) else k[0]
+        assert strategy.k_lower == k_lower
+
+        with pytest.warns(UserWarning):
+            _ = strategy.get_cut_kwargs(self.tape_dags[1])
+
+    @pytest.mark.parametrize("max_wires_by_fragment", [None, [2, 3]])
+    @pytest.mark.parametrize("max_gates_by_fragment", [[20, 30], [20, 30, 40]])
+    def test_by_fragment_sizes(self, max_wires_by_fragment, max_gates_by_fragment):
+        """Test that the user provided by-fragment limits properly propagates."""
+        strategy = qcut.CutStrategy(
+            min_free_wires=2,
+        )
+        if (
+            max_wires_by_fragment
+            and max_gates_by_fragment
+            and len(max_wires_by_fragment) != len(max_gates_by_fragment)
+        ):
+            with pytest.raises(ValueError):
+                cut_kwargs = strategy.get_cut_kwargs(
+                    self.tape_dags[1],
+                    max_wires_by_fragment=max_wires_by_fragment,
+                    max_gates_by_fragment=max_gates_by_fragment,
+                )
+            return
+
+        cut_kwargs = strategy.get_cut_kwargs(
+            self.tape_dags[1],
+            max_wires_by_fragment=max_wires_by_fragment,
+            max_gates_by_fragment=max_gates_by_fragment,
+        )
+        assert len(cut_kwargs) == 1
+
+        cut_kwargs = cut_kwargs[0]
+        assert cut_kwargs["num_fragments"] == len(max_wires_by_fragment or max_gates_by_fragment)
+
+    @pytest.mark.parametrize("max_wires_by_fragment", [2, ["a", 3], [2, 3], None])
+    @pytest.mark.parametrize("max_gates_by_fragment", [2, ["b", 30]])
+    def test_validate_fragment_sizes(self, max_wires_by_fragment, max_gates_by_fragment):
+        """Test that the user provided by-fragment limits has the right types."""
+        with pytest.raises(ValueError):
+            _ = qcut.CutStrategy._validate_input(
+                max_wires_by_fragment=max_wires_by_fragment,
+                max_gates_by_fragment=max_gates_by_fragment,
+            )