Repository: PennyLaneAI/pennylane
Pull request #6019 (2024-07-22 12:12:21)

[Bugfix] `qml.BasisRotation` is not jax compatible
### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
      If you've fixed a bug or added code that should be tested, add a test to the
      test directory!

- [x] All new functions and code must be clearly commented and documented.
      If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [x] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `doc/releases/changelog-dev.md` file, summarizing the
      change, and including a link back to the PR.

- [x] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
      We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

------------------------------------------------------------------------------------------------------------

**Context:** `qml.BasisRotation` is was not jit-compatible.

**Description of the Change:** We modified all the numpy arrays to `jax.numpy` and ensure the tests were passing

**Benefits:** The basis rotation is jittable and thus `qjit` compatible.

**Possible Drawbacks:** jax numpy may be slower than basis numpy.

**Related GitHub Issues:** https://github.com/PennyLaneAI/pennylane/issues/6004


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6019?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 [(`3ad15d7`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/3ad15d7ee01b1c07b74d3d69f82c6c1699967076?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`e33bf03`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/e33bf03d41826dbd553324cd27e0d3288e9ecc96?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

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


```diff
@@            Coverage Diff             @@
##           master    #6019      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         427      427              
  Lines       41057    40776     -281     
==========================================
- Hits        40919    40636     -283     
- Misses        138      140       +2     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6019?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).

josh146: Thanks @obliviateandsurrender and @PabloAMC!

@obliviateandsurrender something odd I notice: using this PR, the following works:

```python
dev = qml.device("lightning.qubit", wires=4)

@jax.jit
@qml.qnode(dev)
def f(U):
    for i in range(4):
        qml.Hadamard(i)
    qml.BasisRotation.compute_decomposition(unitary_matrix=U, wires=[0, 1, 2, 3], check=False)
    return qml.expval(qml.PauliZ(0))

from scipy.stats import unitary_group
U = jnp.array(unitary_group.rvs(4))
f(U)
```

But it no longer works if you remove the `compute_decomposition` (I still get an error message `ValueError: Converting a JAX array to a NumPy array not supported when using the JAX JIT.`).


obliviateandsurrender: Hey @josh146! I am not sure if it is a problem with the template cause it works with `default.qubit`.
I suspect the error is coming from the `convert_to_numpy_parameters` transform being used in the `lightning.qubit` for under `_execute_jvp` function. I am able to turn off this call by setting `diff_method=None` in my qnode.
```
dev = qml.device("lightning.qubit", wires=4)

@jax.jit
@qml.qnode(dev, diff_method=None)
def f(U):
    for i in range(4):
        qml.Hadamard(i)
    qml.BasisRotation(unitary_matrix=U, wires=[0, 1, 2, 3], check=False)
    return qml.expval(qml.PauliZ(0))
U = jnp.array(unitary_group.rvs(4))
```

But this leads to another error as below: 🤔
```
>>> f(U)
TypeError: PhaseShift(): incompatible function arguments. The following argument types are supported:
    1. (self: pennylane_lightning.lightning_qubit_ops.StateVectorC128, arg0: List[int], arg1: bool, arg2: List[float]) -> None
    2. (self: pennylane_lightning.lightning_qubit_ops.StateVectorC128, arg0: List[int], arg1: List[bool], arg2: List[int], arg3: bool, arg4: List[float]) -> None

Invoked with: <pennylane_lightning.lightning_qubit_ops.StateVectorC128 object at 0x17a731c70>, [0], False, [Traced<ShapedArray(float64[])>with<DynamicJaxprTrace(level=1/0)>]
```


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index fd63bce1f..d2bc9de96 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -20,6 +20,9 @@
       - Grover operators
   [(#6379)](https://github.com/PennyLaneAI/pennylane/pull/6379)
 
+* `qml.BasisRotation` template is now JIT compatible.
+  [(#6019)](https://github.com/PennyLaneAI/pennylane/pull/6019)
+
 <h3>Breaking changes 💔</h3>
 
 <h3>Deprecations 👋</h3>
diff --git a/pennylane/qchem/givens_decomposition.py b/pennylane/qchem/givens_decomposition.py
index 18645237c..b09d18ff9 100644
--- a/pennylane/qchem/givens_decomposition.py
+++ b/pennylane/qchem/givens_decomposition.py
@@ -15,8 +15,6 @@
 This module contains the functions needed for performing basis transformations defined by a set of fermionic ladder operators.
 """
 
-import numpy as np
-
 import pennylane as qml
 
 
@@ -38,26 +36,72 @@ def _givens_matrix(a, b, left=True, tol=1e-8):
         tol (float): determines tolerance limits for :math:`|a|` and :math:`|b|` under which they are considered as zero.
 
     Returns:
-        np.ndarray (or tensor): Givens rotation matrix
+        tensor_like: Givens rotation matrix
+
+    Raises:
+        TypeError: if a and b have different interfaces
 
     """
-    abs_a, abs_b = np.abs(a), np.abs(b)
-    if abs_a < tol:
-        cosine, sine, phase = 1.0, 0.0, 1.0
-    elif abs_b < tol:
-        cosine, sine, phase = 0.0, 1.0, 1.0
-    else:
-        hypot = np.hypot(abs_a, abs_b)
-        cosine = abs_b / hypot
-        sine = abs_a / hypot
-        phase = 1.0 * b / abs_b * a.conjugate() / abs_a
+
+    abs_a, abs_b = qml.math.abs(a), qml.math.abs(b)
+    interface_a, interface_b = qml.math.get_interface(a), qml.math.get_interface(b)
+
+    if interface_a != interface_b:
+        raise TypeError(
+            f"The interfaces of 'a' and 'b' do not match. Got {interface_a} and {interface_b}."
+        )
+
+    interface = interface_a
+
+    aprod = qml.math.nan_to_num(abs_b * abs_a)
+    hypot = qml.math.hypot(abs_a, abs_b)
+
+    cosine = qml.math.where(abs_b < tol, 0.0, abs_b / hypot)
+    cosine = qml.math.where(abs_a < tol, 1.0, cosine)
+
+    sine = qml.math.where(abs_b < tol, 1.0, abs_a / hypot)
+    sine = qml.math.where(abs_a < tol, 0.0, sine)
+
+    phase = qml.math.where(abs_b < tol, 1.0, (1.0 * b * qml.math.conj(a)) / (aprod + 1e-15))
+    phase = qml.math.where(abs_a < tol, 1.0, phase)
 
     if left:
-        return np.array([[phase * cosine, -sine], [phase * sine, cosine]])
+        return qml.math.array([[phase * cosine, -sine], [phase * sine, cosine]], like=interface)
+
+    return qml.math.array([[phase * sine, cosine], [-phase * cosine, sine]], like=interface)
+
+
+def _set_unitary_matrix(unitary_matrix, index, value, like=None):
+    """Set the values in the ``unitary_matrix`` at the specified index.
+
+    Args:
+        unitary_matrix (tensor_like): unitary being modified
+        index (Tuple[Int | Ellipsis | List[Int]]): index for slicing the unitary
+        value (tensor_like): new values for the specified index
+        like (str): interface for the unitary matrix
+
+    Returns:
+        tensor_like: modified unitary
+
+    Examples:
+        A = np.eye(5)
+        A = _set_unitary_matrix(A, (0, 0), 5)
+        A = _set_unitary_matrix(A, (1, Ellipsis), np.array([1, 2, 3, 4, 5]))
+        A = _set_unitary_matrix(A, (1, [1, 2]), np.array([3, 4]))
+    """
+    if like is None:
+        like = qml.math.get_interface(unitary_matrix)
+
+    if like == "jax":
+        return unitary_matrix.at[index[0], index[1]].set(
+            value, indices_are_sorted=True, unique_indices=True
+        )
 
-    return np.array([[phase * sine, cosine], [-phase * cosine, sine]])
+    unitary_matrix[index[0], index[1]] = value
+    return unitary_matrix
 
 
+# pylint:disable = too-many-branches
 def givens_decomposition(unitary):
     r"""Decompose a unitary into a sequence of Givens rotation gates with phase shifts and a diagonal phase matrix.
 
@@ -110,7 +154,7 @@ def givens_decomposition(unitary):
         unitary (tensor): unitary matrix on which decomposition will be performed
 
     Returns:
-        (np.ndarray, list[(np.ndarray, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices.
+        (tensor_like, list[(tensor_like, tuple)]): diagonal elements of the phase matrix :math:`D` and Givens rotation matrix :math:`T` with their indices
 
     Raises:
         ValueError: if the provided matrix is not square.
@@ -147,8 +191,10 @@ def givens_decomposition(unitary):
                             # Update U = T(N+j-i-1, N+j-i) @ U
 
     """
+    interface = qml.math.get_deep_interface(unitary)
+    unitary = qml.math.copy(unitary) if interface == "jax" else qml.math.toarray(unitary).copy()
+    M, N = qml.math.shape(unitary)
 
-    unitary, (M, N) = qml.math.toarray(unitary).copy(), unitary.shape
     if M != N:
         raise ValueError(f"The unitary matrix should be of shape NxN, got {unitary.shape}")
 
@@ -158,32 +204,41 @@ def givens_decomposition(unitary):
             for j in range(0, i):
                 indices = [i - j - 1, i - j]
                 grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
-                unitary[:, indices] = unitary[:, indices] @ grot_mat.T
-                right_givens.append((grot_mat.conj(), indices))
+                unitary = _set_unitary_matrix(
+                    unitary, (Ellipsis, indices), unitary[:, indices] @ grot_mat.T, like=interface
+                )
+                right_givens.append((qml.math.conj(grot_mat), indices))
         else:
             for j in range(1, i + 1):
                 indices = [N + j - i - 2, N + j - i - 1]
                 grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
-                unitary[indices] = grot_mat @ unitary[indices]
+                unitary = _set_unitary_matrix(
+                    unitary, (indices, Ellipsis), grot_mat @ unitary[indices, :], like=interface
+                )
                 left_givens.append((grot_mat, indices))
 
     nleft_givens = []
     for grot_mat, (i, j) in reversed(left_givens):
-        sphase_mat = np.diag(np.diag(unitary)[[i, j]])
-        decomp_mat = grot_mat.conj().T @ sphase_mat
+        sphase_mat = qml.math.diag(qml.math.diag(unitary)[qml.math.array([i, j])])
+        decomp_mat = qml.math.conj(grot_mat).T @ sphase_mat
         givens_mat = _givens_matrix(*decomp_mat[1, :].T)
         nphase_mat = decomp_mat @ givens_mat.T
 
         # check for T_{m,n}^{-1} x D = D x T.
-        if not np.allclose(nphase_mat @ givens_mat.conj(), decomp_mat):  # pragma: no cover
+        if not qml.math.is_abstract(decomp_mat) and not qml.math.allclose(
+            nphase_mat @ qml.math.conj(givens_mat), decomp_mat
+        ):  # pragma: no cover
             raise ValueError("Failed to shift phase transposition.")
 
-        unitary[i, i], unitary[j, j] = np.diag(nphase_mat)
-        nleft_givens.append((givens_mat.conj(), (i, j)))
+        for diag_idx, diag_val in zip([(i, i), (j, j)], qml.math.diag(nphase_mat)):
+            unitary = _set_unitary_matrix(unitary, diag_idx, diag_val, like=interface)
+        nleft_givens.append((qml.math.conj(givens_mat), (i, j)))
 
-    phases, ordered_rotations = np.diag(unitary), []
+    phases, ordered_rotations = qml.math.diag(unitary), []
     for grot_mat, (i, j) in list(reversed(nleft_givens)) + list(reversed(right_givens)):
-        if not np.all(np.isreal(grot_mat[0, 1]) and np.isreal(grot_mat[1, 1])):  # pragma: no cover
+        if not qml.math.is_abstract(grot_mat) and not qml.math.all(
+            qml.math.isreal(grot_mat[0, 1]) and qml.math.isreal(grot_mat[1, 1])
+        ):  # pragma: no cover
             raise ValueError(f"Incorrect Givens Rotation encountered, {grot_mat}")
         ordered_rotations.append((grot_mat, (i, j)))
 
diff --git a/pennylane/templates/subroutines/basis_rotation.py b/pennylane/templates/subroutines/basis_rotation.py
index 346e8b837..832a06847 100644
--- a/pennylane/templates/subroutines/basis_rotation.py
+++ b/pennylane/templates/subroutines/basis_rotation.py
@@ -15,8 +15,6 @@
 This module contains the template for performing basis transformation defined by a set of fermionic ladder operators.
 """
 
-import numpy as np
-
 import pennylane as qml
 from pennylane.operation import AnyWires, Operation
 from pennylane.qchem.givens_decomposition import givens_decomposition
@@ -110,15 +108,17 @@ class BasisRotation(Operation):
         return cls._primitive.bind(*wires, unitary_matrix, check=check, id=id)
 
     def __init__(self, wires, unitary_matrix, check=False, id=None):
-        M, N = unitary_matrix.shape
+        M, N = qml.math.shape(unitary_matrix)
+
         if M != N:
-            raise ValueError(
-                f"The unitary matrix should be of shape NxN, got {unitary_matrix.shape}"
-            )
+            raise ValueError(f"The unitary matrix should be of shape NxN, got {(M, N)}")
 
         if check:
-            umat = qml.math.toarray(unitary_matrix)
-            if not np.allclose(umat @ umat.conj().T, np.eye(M, dtype=complex), atol=1e-6):
+            if not qml.math.is_abstract(unitary_matrix) and not qml.math.allclose(
+                unitary_matrix @ qml.math.conj(unitary_matrix).T,
+                qml.math.eye(M, dtype=complex),
+                atol=1e-4,
+            ):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
@@ -153,35 +153,39 @@ class BasisRotation(Operation):
             list[.Operator]: decomposition of the operator
         """
 
-        M, N = unitary_matrix.shape
+        M, N = qml.math.shape(unitary_matrix)
         if M != N:
             raise ValueError(
                 f"The unitary matrix should be of shape NxN, got {unitary_matrix.shape}"
             )
 
         if check:
-            umat = qml.math.toarray(unitary_matrix)
-            if not np.allclose(umat @ umat.conj().T, np.eye(M, dtype=complex), atol=1e-4):
+            if not qml.math.is_abstract(unitary_matrix) and not qml.math.allclose(
+                unitary_matrix @ qml.math.conj(unitary_matrix).T,
+                qml.math.eye(M, dtype=complex),
+                atol=1e-4,
+            ):
                 raise ValueError("The provided transformation matrix should be unitary.")
 
         if len(wires) < 2:
             raise ValueError(f"This template requires at least two wires, got {len(wires)}")
 
         op_list = []
+
         phase_list, givens_list = givens_decomposition(unitary_matrix)
 
         for idx, phase in enumerate(phase_list):
-            op_list.append(qml.PhaseShift(np.angle(phase), wires=wires[idx]))
+            op_list.append(qml.PhaseShift(qml.math.angle(phase), wires=wires[idx]))
 
         for grot_mat, indices in givens_list:
-            theta = np.arccos(np.real(grot_mat[1, 1]))
-            phi = np.angle(grot_mat[0, 0])
+            theta = qml.math.arccos(qml.math.real(grot_mat[1, 1]))
+            phi = qml.math.angle(grot_mat[0, 0])
 
             op_list.append(
                 qml.SingleExcitation(2 * theta, wires=[wires[indices[0]], wires[indices[1]]])
             )
 
-            if not np.isclose(phi, 0.0):
+            if qml.math.is_abstract(phi) or not qml.math.isclose(phi, 0.0):
                 op_list.append(qml.PhaseShift(phi, wires=wires[indices[0]]))
 
         return op_list


Test patch:
diff --git a/tests/qchem/test_givens_rotations.py b/tests/qchem/test_givens_rotations.py
index ebb254cc7..c7e76dbe3 100644
--- a/tests/qchem/test_givens_rotations.py
+++ b/tests/qchem/test_givens_rotations.py
@@ -18,8 +18,13 @@ Unit tests for functions needed for performing givens decomposition of a unitary
 import pytest
 from scipy.stats import unitary_group
 
+import pennylane as qml
 from pennylane import numpy as np
-from pennylane.qchem.givens_decomposition import _givens_matrix, givens_decomposition
+from pennylane.qchem.givens_decomposition import (
+    _givens_matrix,
+    _set_unitary_matrix,
+    givens_decomposition,
+)
 
 
 @pytest.mark.parametrize("left", [True, False])
@@ -126,3 +131,67 @@ def test_givens_decomposition_exceptions(unitary_matrix, msg_match):
 
     with pytest.raises(ValueError, match=msg_match):
         givens_decomposition(unitary_matrix)
+
+
+@pytest.mark.jax
+def test_givens_matrix_exceptions():
+    """Test that _givens_matrix throws an exception if the parameters have different interface."""
+    import jax.numpy as jnp
+
+    a = np.array(1.2)
+    b = jnp.array(2.3)
+
+    with pytest.raises(TypeError, match="The interfaces of 'a' and 'b' do not match."):
+        _givens_matrix(a, b)
+
+
+# pylint:disable = too-many-arguments
+@pytest.mark.parametrize(
+    ("jax", "unitary_matrix", "index", "value", "like", "expected_matrix"),
+    [
+        (False, np.array([[1, 0], [0, 1]]), (0, 0), 5, None, np.array([[5, 0], [0, 1]])),
+        (False, np.array([[1, 0], [0, 1]]), (0, 0), 5, "numpy", np.array([[5, 0], [0, 1]])),
+        (
+            False,
+            np.array([[1, 0], [0, 1]]),
+            (0, Ellipsis),
+            [1, 2],
+            None,
+            np.array([[1, 2], [0, 1]]),
+        ),
+        (
+            False,
+            np.array([[1, 0], [0, 1]]),
+            (0, Ellipsis),
+            [1, 2],
+            "numpy",
+            np.array([[1, 2], [0, 1]]),
+        ),
+        (False, np.array([[1, 0], [0, 1]]), (1, [0, 1]), [1, 2], None, np.array([[1, 0], [1, 2]])),
+        (
+            False,
+            np.array([[1, 0], [0, 1]]),
+            (1, [0, 1]),
+            [1, 2],
+            "numpy",
+            np.array([[1, 0], [1, 2]]),
+        ),
+        (True, [[1, 0], [0, 1]], (0, 0), 5, None, [[5, 0], [0, 1]]),
+        (True, [[1, 0], [0, 1]], (0, 0), 5, "jax", [[5, 0], [0, 1]]),
+        (True, [[1, 0], [0, 1]], (0, Ellipsis), [1, 2], None, [[1, 2], [0, 1]]),
+        (True, [[1, 0], [0, 1]], (0, Ellipsis), [1, 2], "jax", [[1, 2], [0, 1]]),
+        (True, [[1, 0], [0, 1]], (1, [0, 1]), [1, 2], None, [[1, 0], [1, 2]]),
+        (True, [[1, 0], [0, 1]], (1, [0, 1]), [1, 2], "jax", [[1, 0], [1, 2]]),
+    ],
+)
+@pytest.mark.jax
+def test_set_unitary_matrix(jax, unitary_matrix, index, value, like, expected_matrix):
+    """Test the _set_unitary function on different interfaces."""
+    import jax.numpy as jnp
+
+    if jax:
+        unitary_matrix = jnp.array(unitary_matrix)
+        expected_matrix = jnp.array(expected_matrix)
+
+    new_unitary_matrix = _set_unitary_matrix(unitary_matrix, index, value, like)
+    assert qml.math.allclose(new_unitary_matrix, expected_matrix)
diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py
index c5c9c6434..055566411 100644
--- a/tests/templates/test_subroutines/test_basis_rotation.py
+++ b/tests/templates/test_subroutines/test_basis_rotation.py
@@ -334,11 +334,12 @@ class TestInputs:
         assert template.id == "a"
 
 
-def circuit_template(unitary_matrix):
-    qml.BasisState(np.array([1, 1, 0]), wires=[0, 1, 2])
+def circuit_template(unitary_matrix, check=False):
+    qml.BasisState(qml.math.array([1, 1, 0]), wires=[0, 1, 2])
     qml.BasisRotation(
         wires=range(3),
         unitary_matrix=unitary_matrix,
+        check=check,
     )
     return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
@@ -402,7 +403,7 @@ class TestInterfaces:
         assert np.allclose(grads, np.zeros_like(unitary_matrix, dtype=complex), atol=tol, rtol=0)
 
     @pytest.mark.jax
-    def test_jax(self, tol):
+    def test_jax_jit(self, tol):
         """Test the jax interface."""
 
         import jax
@@ -415,36 +416,26 @@ class TestInterfaces:
                 [-0.58608928 + 0.0j, 0.03902657 + 0.04633548j, -0.57220635 + 0.57044649j],
             ]
         )
-        weights = jnp.array(
-            [
-                2.2707802713289267,
-                2.9355948424220206,
-                -1.4869222527726533,
-                1.2601662579297865,
-                2.3559705032936717,
-                1.1748572730890159,
-                2.2500537657656356,
-                -0.7251404204443089,
-                2.3577346350335198,
-            ]
-        )
 
         dev = qml.device("default.qubit", wires=3)
 
-        circuit = qml.QNode(circuit_template, dev)
-        circuit2 = qml.QNode(circuit_decomposed, dev)
+        circuit = jax.jit(qml.QNode(circuit_template, dev), static_argnames="check")
+        circuit2 = qml.QNode(circuit_template, dev)
 
         res = circuit(unitary_matrix)
-        res2 = circuit2(weights)
+        res2 = circuit2(unitary_matrix)
+        res3 = circuit2(qml.math.toarray(unitary_matrix))
+
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+        assert qml.math.allclose(res, res3, atol=tol, rtol=0)
 
         grad_fn = jax.grad(circuit)
         grads = grad_fn(unitary_matrix)
 
         grad_fn2 = jax.grad(circuit2)
-        grads2 = grad_fn2(weights)
+        grads2 = grad_fn2(unitary_matrix)
 
-        assert np.allclose(grads[0], grads2[0], atol=tol, rtol=0)
+        assert qml.math.allclose(grads, grads2, atol=tol, rtol=0)
 
     @pytest.mark.tf
     def test_tf(self, tol):