Repository: PennyLaneAI/pennylane
Pull request #5972 (2024-07-09 04:44:56)

Add exponential fitting for ZNE
**Context:** As part of increasing the utility of ZNE within Catalyst, we'd like to add the option to use exponential fitting with ZNE. The issue 

**Description of the Change:** A function which is written using JAX is added which performs an exponential fit of data, and evaluates said fit at 0.

**Benefits:** Increased capabilities for users to play with. This is especially important given the benefits of tuning ZNE to the problem and backend at hand.

**Possible Drawbacks:** The current implementation transforms data using the natural logarithm $\ln$ to linearize the data so that a linear fit can be applied to learn the necessary parameters (see section below on the exponential model). Once they transformed parameters are learned, they can be converted back into those from the exponential model quite simply. The major drawback of this approach is that they expectation values have to be shifted to be in the domain of $\ln$ (i.e. $> 0$), or thrown away if small enough (see $\varepsilon$ `eps` parameter introduced in `expontial_extrapolate`).

**Related GitHub Issues:** https://github.com/PennyLaneAI/catalyst/issues/754

**Exponential model:** Assuming an exponential model with asymptote $y(x) = A\mathrm{e}^{Bx} + C$ we must first turn this into a linear equation to be able to apply a linear fit.

$$y(x) = A\mathrm{e}^{Bx} + C$$

$$y(x) - C = A\mathrm{e}^{Bx}$$

$$\ln(y(x) - C) = \ln(A) + Bx$$

$$\tilde{y}(x) = \tilde{A} + Bx$$

Once $\tilde{A}$ is learned, the hard part is over as the ZNE value is "simply" $y(0) = A + C$.

josh146: Thanks for this PR @natestemen!

> ⚠️ draft PR note: I'm currently having trouble getting tests to pass on this change. I am aware and actively working on it. Putting this PR up for visibility and in case anyone has clues where I'm messing things up.

Feel free to tag us directly on the PR, or alternatively let us know if you want anyone to go through the PR.
natestemen: Guidance would be appreciated on resolving the CI error raised [here](https://github.com/PennyLaneAI/pennylane/actions/runs/9903181919/job/27358295670?pr=5972#step:8:1302).
natestemen: @josh146 ready for the first look!
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5972?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.65%. Comparing base [(`94b5533`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/94b5533a0d6a7b0653f66050e9e514fdd7d9f0e0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`169d53e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/169d53e8c680572e8bc1dd420d0e5296c71174bb?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 2 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5972      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         427      427              
  Lines       40983    40706     -277     
==========================================
- Hits        40845    40567     -278     
- Misses        138      139       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5972?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

EmilianoG-byte: Is there a reason why we are not checking differentiability for other backends besides `jax`? For instance here:

https://github.com/PennyLaneAI/pennylane/blob/169d53e8c680572e8bc1dd420d0e5296c71174bb/tests/transforms/test_mitigate.py#L646-L647

Why didn't we add

```python
@pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
```

like we did with JAX and JAX-JIT?
natestemen: Thanks for the docstring suggestions! Addressed in https://github.com/PennyLaneAI/pennylane/pull/5972/commits/e0815dac01fd235d8177b2f36d45e0c927cf284f.

> Is there a reason why we are not checking differentiability for other backends besides jax?

The goal of this PR was to get an exponential fit function in that would be compatible with [`pennylaneai/catalyst`](https://github.com/pennylaneai/catalyst). Originally I thought the `exponential_extrapolate` function was required to be written in JAX and hence only ensured it was tested as such. I wasn't aware of the testing requirements of the other backends[^1] until getting some review on this PR.

Happy to add your suggestion. Let me know.

[^1]: I'm still unclear if the other backends (torch, tf, autograd) are required in order for this functionality to work with catalyst. Is it?
EmilianoG-byte: > 1. I'm still unclear if the other backends (torch, tf, autograd) are required in order for this functionality to work with catalyst. Is it? [↩](#user-content-fnref-1-f78a22c9d9339f481c026149f1235a68)

Good question @natestemen 🤔 

Strictly speaking, the other interfaces are not needed for this functionality to work with Catalyst. However, since this will be part of the Pennylane features (even without Jitting our workflows), it should be compatible with the interfaces supported by Pennylane - this includes being able to differentiate it. 

The good thing is that since you are using `qml.math` in `exponential_extrapolate`, adding this support should be rather straightforward by parametrizing the differentiability tests with 

```python 
@pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
```

wherever only `richardson_extrapolate` was used previously. Let me know if there is any unforeseen difficulty 😃.


rmoyard: @natestemen The requirement for it to work with Jax is to be Jax-jittable. I have asked for other interfaces because in PL we always make things compatible for the different frameworks.

Patch:
diff --git a/doc/development/deprecations.rst b/doc/development/deprecations.rst
index b45525a08..a1ef7ce72 100644
--- a/doc/development/deprecations.rst
+++ b/doc/development/deprecations.rst
@@ -9,6 +9,16 @@ deprecations are listed below.
 Pending deprecations
 --------------------
 
+* The ``decomp_depth`` argument in ``qml.device`` is deprecated. 
+
+  - Deprecated in v0.38
+  - Will be removed in v0.39
+
+* The ``max_expansion`` argument in ``qml.QNode`` is deprecated. 
+
+  - Deprecated in v0.38
+  - Will be removed in v0.39
+
 * The functions ``qml.transforms.sum_expand`` and ``qml.transforms.hamiltonian_expand`` are deprecated.
   Instead, ``qml.transforms.split_non_commuting`` can be used for equivalent behaviour.
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b4e808eb0..c7bceecd5 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -22,11 +22,15 @@
 
 * `SProd.terms` now flattens out the terms if the base is a multi-term observable.
   [(#5885)](https://github.com/PennyLaneAI/pennylane/pull/5885)
-  
+
 * A new method `to_mat` has been added to the `FermiWord` and `FermiSentence` classes, which allows
   computing the matrix representation of these Fermi operators.
   [(#5920)](https://github.com/PennyLaneAI/pennylane/pull/5920)
 
+* New functionality has been added to natively support exponential extrapolation when using the `mitigate_with_zne`. This allows
+  users to have more control over the error mitigation protocol without needing to add further dependencies.
+  [(#5972)](https://github.com/PennyLaneAI/pennylane/pull/5972)
+
 <h3>Improvements 🛠</h3>
 
 * `StateMP.process_state` defines rules in `cast_to_complex` for complex casting, avoiding a superfluous state vector copy in Lightning simulations
@@ -87,6 +91,12 @@
 
 <h3>Deprecations 👋</h3>
 
+* The `decomp_depth` argument in `qml.device` has been deprecated.
+  [(#6026)](https://github.com/PennyLaneAI/pennylane/pull/6026)
+
+* The `max_expansion` argument in `qml.QNode` has been deprecated.
+  [(#6026)](https://github.com/PennyLaneAI/pennylane/pull/6026)
+
 * The `expansion_strategy` attribute in the `QNode` class is deprecated.
   [(#5989)](https://github.com/PennyLaneAI/pennylane/pull/5989)
 
@@ -169,4 +179,5 @@ Christina Lee,
 William Maxwell,
 Vincent Michaud-Rioux,
 Mudit Pandey,
-Erik Schultheis.
+Erik Schultheis,
+nate stemen.
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 04b7e4f24..58cf18f38 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -18,6 +18,7 @@ PennyLane can be directly imported.
 from importlib import reload, metadata
 from sys import version_info
 
+import warnings
 import numpy as _np
 
 from semantic_version import SimpleSpec, Version
@@ -262,6 +263,10 @@ def device(name, *args, **kwargs):
         decomp_depth (int): For when custom decompositions are specified,
             the maximum expansion depth used by the expansion function.
 
+    .. warning::
+
+        The ``decomp_depth`` argument is deprecated and will be removed in version 0.39.
+
     All devices must be loaded by specifying their **short-name** as listed above,
     followed by the **wires** (subsystems) you wish to initialize. The ``wires``
     argument can be an integer, in which case the wires of the device are addressed
@@ -393,7 +398,15 @@ def device(name, *args, **kwargs):
         # Pop the custom decomposition keyword argument; we will use it here
         # only and not pass it to the device.
         custom_decomps = kwargs.pop("custom_decomps", None)
-        decomp_depth = kwargs.pop("decomp_depth", 10)
+        decomp_depth = kwargs.pop("decomp_depth", None)
+
+        if decomp_depth is not None:
+            warnings.warn(
+                "The decomp_depth argument is deprecated and will be removed in version 0.39. ",
+                PennyLaneDeprecationWarning,
+            )
+        else:
+            decomp_depth = 10
 
         kwargs.pop("config", None)
         options.update(kwargs)
@@ -414,6 +427,7 @@ def device(name, *args, **kwargs):
 
         # Once the device is constructed, we set its custom expansion function if
         # any custom decompositions were specified.
+
         if custom_decomps is not None:
             if isinstance(dev, pennylane.devices.LegacyDevice):
                 custom_decomp_expand_fn = pennylane.transforms.create_decomp_expand_fn(
diff --git a/pennylane/capture/capture_qnode.py b/pennylane/capture/capture_qnode.py
index d0faf65b6..0b7ecb676 100644
--- a/pennylane/capture/capture_qnode.py
+++ b/pennylane/capture/capture_qnode.py
@@ -15,6 +15,7 @@
 This submodule defines a capture compatible call to QNodes.
 """
 
+import warnings
 from copy import copy
 from dataclasses import asdict
 from functools import lru_cache, partial
@@ -29,7 +30,7 @@ except ImportError:
 
 
 def _get_shapes_for(*measurements, shots=None, num_device_wires=0):
-    if jax.config.jax_enable_x64:
+    if jax.config.jax_enable_x64:  # pylint: disable=no-member
         dtype_map = {
             float: jax.numpy.float64,
             int: jax.numpy.int64,
@@ -65,7 +66,13 @@ def _get_qnode_prim():
         def qfunc(*inner_args):
             return jax.core.eval_jaxpr(qfunc_jaxpr.jaxpr, qfunc_jaxpr.consts, *inner_args)
 
-        qnode = qml.QNode(qfunc, device, **qnode_kwargs)
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                action="ignore",
+                message=r"The max_expansion argument is deprecated and will be removed in version 0.39.",
+                category=qml.PennyLaneDeprecationWarning,
+            )
+            qnode = qml.QNode(qfunc, device, **qnode_kwargs)
         return qnode._impl_call(*args, shots=shots)  # pylint: disable=protected-access
 
     # pylint: disable=unused-argument
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 83dcb1ae5..fbc3dff55 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -93,6 +93,7 @@ Transforms for error mitigation
     ~transforms.fold_global
     ~transforms.poly_extrapolate
     ~transforms.richardson_extrapolate
+    ~transforms.exponential_extrapolate
 
 Other transforms
 ~~~~~~~~~~~~~~~~
@@ -322,7 +323,13 @@ from .split_non_commuting import split_non_commuting
 from .split_to_single_terms import split_to_single_terms
 from .insert_ops import insert
 
-from .mitigate import mitigate_with_zne, fold_global, poly_extrapolate, richardson_extrapolate
+from .mitigate import (
+    mitigate_with_zne,
+    fold_global,
+    poly_extrapolate,
+    richardson_extrapolate,
+    exponential_extrapolate,
+)
 from .optimization import (
     cancel_inverses,
     commute_controlled,
diff --git a/pennylane/transforms/mitigate.py b/pennylane/transforms/mitigate.py
index 572c904bc..59d47e06c 100644
--- a/pennylane/transforms/mitigate.py
+++ b/pennylane/transforms/mitigate.py
@@ -318,6 +318,48 @@ def richardson_extrapolate(x, y):
     return poly_extrapolate(x, y, len(x) - 1)
 
 
+def exponential_extrapolate(x, y, asymptote=None, eps=1.0e-6):
+    r"""Extrapolate to the zero-noise limit using an exponential model (:math:`Ae^{Bx} + C`). This
+    is done by linearizing the data using a logarithm, whereupon a linear fit is performed. Once
+    the model parameters are found, they are transformed back to exponential parameters.
+
+    Args:
+        x (Array): Data in x axis.
+        y (Array): Data in y axis such that :math:`y = f(x)`.
+        asymptote (float): Infinite noise limit expected for your circuit of interest (:math:`C`
+            in the equation above). Defaults to 0 in the case an asymptote is not supplied.
+        eps (float): Epsilon to regularize :math:`\log(y - C)` when the argument is to close to
+            zero or negative.
+
+    Returns:
+        float: Extrapolated value at f(0).
+
+    .. seealso:: :func:`~.pennylane.transforms.richardson_extrapolate`, :func:`~.pennylane.transforms.mitigate_with_zne`.
+
+    **Example:**
+
+    >>> np.random.seed(0)
+    >>> x = np.linspace(1, 10, 5)
+    >>> y = np.exp(-x) + np.random.normal(scale=0.1, size=len(x))
+    >>> qml.transforms.exponential_extrapolate(x, y)
+    0.23365009000522544
+    """
+    y = qml.math.stack(y)
+    slope, y_intercept = _polyfit(x, y, 1)
+    if asymptote is None:
+        sign = qml.math.sign(-slope)
+        asymptote = 0.0
+    else:
+        sign = qml.math.sign(-(asymptote - y_intercept))
+
+    y_shifted = sign * (y - asymptote)
+    y_shifted = qml.math.where(y_shifted < eps, eps, y_shifted)
+    y_scaled = qml.math.log(y_shifted)
+
+    zne_unscaled = poly_extrapolate(x, y_scaled, 1)
+    return sign * qml.math.exp(zne_unscaled) + asymptote
+
+
 # pylint: disable=too-many-arguments, protected-access
 @transform
 def mitigate_with_zne(
diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index 7dab18e3c..09c667bdd 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -239,7 +239,12 @@ class QNode:
             on supported options for your chosen gradient transform.
 
     .. warning::
-        The ``expansion_strategy`` attribute is deprecated and will be removed in version 0.39.
+
+        The ``expansion_strategy`` argument is deprecated and will be removed in version 0.39.
+
+    .. warning::
+
+        The ``max_expansion`` argument is deprecated and will be removed in version 0.39.
 
     **Example**
 
@@ -452,7 +457,7 @@ class QNode:
         interface="auto",
         diff_method="best",
         expansion_strategy=None,
-        max_expansion=10,
+        max_expansion=None,
         grad_on_execution="best",
         cache="auto",
         cachesize=10000,
@@ -462,6 +467,15 @@ class QNode:
         mcm_method=None,
         **gradient_kwargs,
     ):
+        # Moving it here since the old default value is checked on debugging
+        if max_expansion is not None:
+            warnings.warn(
+                "The max_expansion argument is deprecated and will be removed in version 0.39. ",
+                qml.PennyLaneDeprecationWarning,
+            )
+        else:
+            max_expansion = 10
+
         if logger.isEnabledFor(logging.DEBUG):
             logger.debug(
                 """Creating QNode(func=%s, device=%s, interface=%s, diff_method=%s, expansion_strategy=%s, max_expansion=%s, grad_on_execution=%s, cache=%s, cachesize=%s, max_diff=%s, gradient_kwargs=%s""",


Test patch:
diff --git a/tests/test_device.py b/tests/test_device.py
index 1639f76ed..2c15a7992 100644
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -1069,6 +1069,15 @@ class TestDeviceInit:
 
         assert dev.shots == 22
 
+    def test_decomp_depth_is_deprecated(self):
+        """Test that a warning is raised when using the deprecated decomp_depth argument"""
+
+        with pytest.warns(
+            qml.PennyLaneDeprecationWarning,
+            match="The decomp_depth argument is deprecated",
+        ):
+            qml.device("default.qubit", decomp_depth=1)
+
 
 class TestBatchExecution:
     """Tests for the batch_execute method."""
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index 71e99e705..3251a7250 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -101,6 +101,15 @@ class TestInitialization:
 
         assert f.execute_kwargs["cache"] is True
 
+    def test_max_expansion_is_deprecated(self):
+        """Test that a warning is raised when using the deprecated max_expansion argument"""
+        dev = qml.device("default.qubit", wires=1)
+        with pytest.warns(
+            qml.PennyLaneDeprecationWarning,
+            match="The max_expansion argument is deprecated",
+        ):
+            QNode(dummyfunc, dev, max_expansion=10)
+
 
 # pylint: disable=too-many-public-methods
 class TestValidation:
diff --git a/tests/transforms/test_mitigate.py b/tests/transforms/test_mitigate.py
index 7cf6dd179..002a7b1af 100644
--- a/tests/transforms/test_mitigate.py
+++ b/tests/transforms/test_mitigate.py
@@ -23,7 +23,12 @@ from packaging import version
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.tape import QuantumScript
-from pennylane.transforms import fold_global, mitigate_with_zne, richardson_extrapolate
+from pennylane.transforms import (
+    exponential_extrapolate,
+    fold_global,
+    mitigate_with_zne,
+    richardson_extrapolate,
+)
 
 with qml.queuing.AnnotatedQueue() as q_tape:
     qml.BasisState([1], wires=0)
@@ -109,7 +114,8 @@ class TestMitigateWithZNE:
         for t in tapes:
             same_tape(t, tape)
 
-    def test_multi_returns(self):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_multi_returns(self, extrapolate):
         """Tests if the expected shape is returned when mitigating a circuit with two returns"""
         noise_strength = 0.05
 
@@ -126,7 +132,7 @@ class TestMitigateWithZNE:
             qml.transforms.mitigate_with_zne,
             scale_factors=[1, 2, 3],
             folding=fold_global,
-            extrapolate=richardson_extrapolate,
+            extrapolate=extrapolate,
         )
         @qml.qnode(dev)
         def mitigated_circuit(w1, w2):
@@ -507,17 +513,67 @@ class TestDifferentiableZNE:
         coeffs = qml.transforms.mitigate._polyfit(x, y, 2)
         assert qml.math.allclose(qml.math.squeeze(coeffs), [3, 2, 1])
 
+    @pytest.mark.parametrize("exp_params", [[0.5, -2, 2], [-9, -4, 0]])
+    def test_exponential_extrapolation_accuracy(self, exp_params):
+        """Testing the exponential extrapolation works as expected for known exponential models."""
+        A, B, asymptote = exp_params
+        x = np.linspace(1, 4, 4)
+        y = A * np.exp(B * x) + asymptote
+        zne_val = qml.transforms.exponential_extrapolate(x, y, asymptote=asymptote)
+        assert qml.math.allclose(zne_val, A + asymptote, atol=1e-3)
+
+    @pytest.mark.autograd
+    def test_exponential_extrapolation_autograd(self):
+        """Test exponential extrapolation works with expvals stored as a numpy array."""
+        scale_factors = [1, 3, 5]
+        noise_scaled_expvals = np.array([0.9, 0.8, 0.7])
+        zne_val = qml.transforms.exponential_extrapolate(scale_factors, noise_scaled_expvals)
+        assert isinstance(zne_val, np.ndarray)
+        assert zne_val.ndim == 0
+
+    @pytest.mark.tf
+    def test_exponential_extrapolation_tf(self):
+        """Test exponential extrapolation works with expvals stored as a tensorflow tensor."""
+        import tensorflow as tf
+
+        scale_factors = [1, 3, 5]
+        noise_scaled_expvals = tf.constant([0.9, 0.8, 0.7], dtype=tf.float32)
+        zne_val = qml.transforms.exponential_extrapolate(scale_factors, noise_scaled_expvals)
+        assert tf.is_tensor(zne_val)
+        assert zne_val.shape.ndims == 0
+
+    @pytest.mark.torch
+    def test_exponential_extrapolation_torch(self):
+        """Test exponential extrapolation works with expvals stored as a torch tensor."""
+        import torch
+
+        scale_factors = [1, 3, 5]
+        noise_scaled_expvals = torch.tensor([0.9, 0.8, 0.7])
+        zne_val = qml.transforms.exponential_extrapolate(scale_factors, noise_scaled_expvals)
+        assert torch.is_tensor(zne_val)
+        assert zne_val.ndimension() == 0
+
+    @pytest.mark.jax
+    def test_exponential_extrapolation_jax(self):
+        """Test exponential extrapolation works with expvals stored as a jax array."""
+        import jax.numpy as jnp
+
+        scale_factors = [1, 3, 5]
+        noise_scaled_expvals = jnp.array([0.9, 0.8, 0.7])
+        zne_val = qml.transforms.exponential_extrapolate(scale_factors, noise_scaled_expvals)
+        assert isinstance(zne_val, jnp.ndarray)
+        assert zne_val.ndim == 0
+
     @pytest.mark.autograd
-    def test_diffability_autograd(self):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_autograd(self, extrapolate):
         """Testing that the mitigated qnode can be differentiated and returns the correct gradient in autograd"""
         qnode_noisy = qml.QNode(qfunc, dev_noisy)
         qnode_ideal = qml.QNode(qfunc, dev_ideal)
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(
-            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
-        )
+        mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
 
         theta = np.array([np.pi / 4, np.pi / 4], requires_grad=True)
 
@@ -531,7 +587,8 @@ class TestDifferentiableZNE:
 
     @pytest.mark.jax
     @pytest.mark.parametrize("interface", ["auto", "jax"])
-    def test_diffability_jax(self, interface):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_jax(self, interface, extrapolate):
         """Testing that the mitigated qnode can be differentiated and returns the correct gradient in jax"""
         import jax
         import jax.numpy as jnp
@@ -541,9 +598,7 @@ class TestDifferentiableZNE:
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(
-            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
-        )
+        mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
 
         theta = jnp.array(
             [np.pi / 4, np.pi / 4],
@@ -559,7 +614,8 @@ class TestDifferentiableZNE:
 
     @pytest.mark.jax
     @pytest.mark.parametrize("interface", ["auto", "jax", "jax-jit"])
-    def test_diffability_jaxjit(self, interface):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_jaxjit(self, interface, extrapolate):
         """Testing that the mitigated qnode can be differentiated and returns the correct gradient in jax-jit"""
         import jax
         import jax.numpy as jnp
@@ -570,7 +626,7 @@ class TestDifferentiableZNE:
         scale_factors = [1.0, 2.0, 3.0]
 
         mitigated_qnode = jax.jit(
-            mitigate_with_zne(qnode_noisy, scale_factors, fold_global, richardson_extrapolate)
+            mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
         )
 
         theta = jnp.array(
@@ -587,7 +643,8 @@ class TestDifferentiableZNE:
 
     @pytest.mark.torch
     @pytest.mark.parametrize("interface", ["auto", "torch"])
-    def test_diffability_torch(self, interface):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_torch(self, interface, extrapolate):
         """Testing that the mitigated qnode can be differentiated and returns the correct gradient in torch"""
         import torch
 
@@ -596,9 +653,7 @@ class TestDifferentiableZNE:
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(
-            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
-        )
+        mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
 
         theta = torch.tensor([np.pi / 4, np.pi / 4], requires_grad=True)
 
@@ -616,7 +671,8 @@ class TestDifferentiableZNE:
 
     @pytest.mark.tf
     @pytest.mark.parametrize("interface", ["auto", "tf"])
-    def test_diffability_tf(self, interface):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_tf(self, interface, extrapolate):
         """Testing that the mitigated qnode can be differentiated and returns the correct gradient in tf"""
         import tensorflow as tf
 
@@ -625,9 +681,7 @@ class TestDifferentiableZNE:
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(
-            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
-        )
+        mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
 
         theta = tf.Variable([np.pi / 4, np.pi / 4])
 
@@ -645,7 +699,8 @@ class TestDifferentiableZNE:
         assert qml.math.allclose(grad, grad_ideal, atol=1e-2)
 
     @pytest.mark.autograd
-    def test_diffability_autograd_multi(self):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_autograd_multi(self, extrapolate):
         """Testing that the mitigated qnode can be differentiated and returns
         the correct gradient in autograd for multiple measurements"""
         qnode_noisy = qml.QNode(qfunc_multi, dev_noisy)
@@ -653,9 +708,7 @@ class TestDifferentiableZNE:
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(
-            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
-        )
+        mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
 
         theta = np.array([np.pi / 4, np.pi / 6], requires_grad=True)
 
@@ -669,7 +722,8 @@ class TestDifferentiableZNE:
 
     @pytest.mark.jax
     @pytest.mark.parametrize("interface", ["auto", "jax"])
-    def test_diffability_jax_multi(self, interface):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_jax_multi(self, interface, extrapolate):
         """Testing that the mitigated qnode can be differentiated and returns
         the correct gradient in jax for multiple measurements"""
         import jax
@@ -680,9 +734,7 @@ class TestDifferentiableZNE:
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(
-            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
-        )
+        mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
 
         theta = jnp.array(
             [np.pi / 4, np.pi / 6],
@@ -698,7 +750,8 @@ class TestDifferentiableZNE:
 
     @pytest.mark.jax
     @pytest.mark.parametrize("interface", ["auto", "jax", "jax-jit"])
-    def test_diffability_jaxjit_multi(self, interface):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_jaxjit_multi(self, interface, extrapolate):
         """Testing that the mitigated qnode can be differentiated and
         returns the correct gradient in jax-jit for multiple measurements"""
         import jax
@@ -710,7 +763,7 @@ class TestDifferentiableZNE:
         scale_factors = [1.0, 2.0, 3.0]
 
         mitigated_qnode = jax.jit(
-            mitigate_with_zne(qnode_noisy, scale_factors, fold_global, richardson_extrapolate)
+            mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
         )
 
         theta = jnp.array(
@@ -727,7 +780,8 @@ class TestDifferentiableZNE:
 
     @pytest.mark.torch
     @pytest.mark.parametrize("interface", ["auto", "torch"])
-    def test_diffability_torch_multi(self, interface):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_torch_multi(self, interface, extrapolate):
         """Testing that the mitigated qnode can be differentiated and returns
         the correct gradient in torch for multiple measurements"""
         import torch
@@ -737,9 +791,7 @@ class TestDifferentiableZNE:
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(
-            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
-        )
+        mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
 
         theta = torch.tensor([np.pi / 4, np.pi / 6], requires_grad=True)
 
@@ -756,7 +808,8 @@ class TestDifferentiableZNE:
         assert qml.math.allclose(grad, grad_ideal, atol=1e-2)
 
     @pytest.mark.tf
-    def test_diffability_tf_multi(self):
+    @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])
+    def test_diffability_tf_multi(self, extrapolate):
         """Testing that the mitigated qnode can be differentiated and returns
         the correct gradient in tf for multiple measurements"""
         import tensorflow as tf
@@ -766,9 +819,7 @@ class TestDifferentiableZNE:
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(
-            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
-        )
+        mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate)
 
         theta = tf.Variable([np.pi / 4, np.pi / 6])
 
diff --git a/tests/transforms/test_tape_expand.py b/tests/transforms/test_tape_expand.py
index 56fee6948..5d7b59a67 100644
--- a/tests/transforms/test_tape_expand.py
+++ b/tests/transforms/test_tape_expand.py
@@ -17,6 +17,8 @@ Unit tests for tape expansion stopping criteria and expansion functions.
 # pylint: disable=too-few-public-methods, invalid-unary-operand-type, no-member,
 # pylint: disable=arguments-differ, arguments-renamed,
 
+import warnings
+
 import numpy as np
 import pytest
 
@@ -24,15 +26,6 @@ import pennylane as qml
 from pennylane.wires import Wires
 
 
-@pytest.fixture(scope="function", autouse=True)
-def capture_warnings(recwarn):
-    yield
-    if len(recwarn) > 0:
-        for w in recwarn:
-            assert isinstance(w.message, qml.PennyLaneDeprecationWarning)
-            assert "'expansion_strategy' attribute is deprecated" in str(w.message)
-
-
 class TestCreateExpandFn:
     """Test creating expansion functions from stopping criteria."""
 
@@ -438,6 +431,21 @@ def custom_basic_entangler_layers(weights, wires, **kwargs):
 class TestCreateCustomDecompExpandFn:
     """Tests for the custom_decomps argument for devices"""
 
+    @pytest.fixture(scope="function", autouse=True)
+    def capture_warnings(self):
+        with pytest.warns(
+            qml.PennyLaneDeprecationWarning,
+        ) as record:
+            yield
+
+            assert any(
+                "'expansion_strategy' attribute is deprecated" in str(w.message) for w in record
+            )
+
+        for w in record:
+            if "'expansion_strategy' attribute is deprecated" not in str(w.message):
+                warnings.warn(w.message, w.category)
+
     @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
     def test_string_and_operator_allowed(self, device_name):
         """Test that the custom_decomps dictionary accepts both strings and operator classes as keys."""
@@ -544,7 +552,12 @@ class TestCreateCustomDecompExpandFn:
         """Test that specifying a single custom decomposition works as expected."""
 
         custom_decomps = {"Hadamard": custom_hadamard, "CNOT": custom_cnot}
-        decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=0)
+        with pytest.warns(
+            qml.PennyLaneDeprecationWarning, match="The decomp_depth argument is deprecated"
+        ):
+            decomp_dev = qml.device(
+                device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=0
+            )
 
         @qml.qnode(decomp_dev, expansion_strategy="device")
         def circuit():
@@ -721,13 +734,16 @@ class TestCreateCustomDecompExpandFn:
         # not be further decomposed even though the custom decomposition is specified.
         custom_decomps = {"BasicEntanglerLayers": custom_basic_entangler_layers, "RX": custom_rx}
 
-        decomp_dev_2 = qml.device(
-            device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=2
-        )
+        with pytest.warns(
+            qml.PennyLaneDeprecationWarning, match="The decomp_depth argument is deprecated"
+        ):
+            decomp_dev_2 = qml.device(
+                device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=2
+            )
 
-        decomp_dev_3 = qml.device(
-            device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=3
-        )
+            decomp_dev_3 = qml.device(
+                device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=3
+            )
 
         def circuit():
             qml.BasicEntanglerLayers([[0.1, 0.2]], wires=[0, 1])
@@ -910,7 +926,7 @@ class TestCreateCustomDecompExpandFn:
             custom_decomps = {"MultiRZ": qml.MultiRZ.compute_decomposition}
             dev = qml.device("lightning.qubit", wires=2, custom_decomps=custom_decomps)
 
-            @qml.qnode(dev, diff_method="adjoint")
+            @qml.qnode(dev, diff_method="adjoint", expansion_strategy="device")
             def cost(theta):
                 qml.Hadamard(wires=0)
                 qml.Hadamard(wires=1)