Repository: PennyLaneAI/pennylane (python)

Issue:
Multi dispatch `tensordot`
### Feature details

It would be nice to be able to do
```python
qml.math.tensordot(a_list, an_autograd_array, axes=axes)
```
with `tensordot` dispatching to Autograd. Currently, the value of the above operation is correct but its derivative is `0`:
```python
x = qml.math.array([0.2, 0.1])
a_list = [1.2, 0.8]

f = lambda x: qml.math.tensordot(a_list, x, axes=[[0], [0]])
f(x) # Evaluates correctly to 0.32
qml.grad(f)(x) # Evaluates wrongly to [0.0, 0.0]
```
When using the Autograd array as first entry, the derivative is correct:

```python
g = lambda x: qml.math.tensordot(x, a_list, axes=[[0], [0]])
g(x) # Evaluates correctly to 0.32
qml.grad(g)(x) # Evaluates correctly to [1.2, 0.8]
``` 

Therefore, a multi dispatch over the arguments of `tensordot` would be nice to have.

### Implementation

One could probably implement this quite similar to #1816 

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

_No response_

Pull request:
Use the new `multi_dispatch` decorator in the math module. Add `tensordot` tests.
**Context:**
This PR makes use of the `multi_dispatch` decorator that was introduced in #2084 throughout `pennylane/math/multi_dispatch.py`. It also introduces tests of differentiability in all interfaces for `qml.math.tensordot`, as well as a minimal test of `tensordot` with Sequences as inputs.

#1865 was partially closed by #1725, but a `qml.math.coerce` is introduced here to `tensordot` to complete this patch of `tensordot`.

**Description of the Change:**
Some more tests; Usage of `multi_dispatch`.

**Benefits:**
Concise codebase and usage examples of the new decorator; Test coverage

**Possible Drawbacks:**
Tests take longer.
Performance might be reduced minimally by using a higher-level object (the decorator) rather than hard-coding the usage of `_multi_dispatch` to determine the interface.

**Related GitHub Issues:**
closes #1865 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 097ac57e4..830471f99 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -327,10 +327,13 @@
   [(#2063)](https://github.com/PennyLaneAI/pennylane/pull/2063)
 
 * Added a new `multi_dispatch` decorator that helps ease the definition of new functions
-  inside PennyLane. We can decorate the function, indicating the arguments that are
-  tensors handled by the interface:
+  inside PennyLane. The decorator is used throughout the math module, demonstrating use cases.
   [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084)
-    
+  [(#2096)](https://github.com/PennyLaneAI/pennylane/pull/2096)
+
+  We can decorate a function, indicating the arguments that are
+  tensors handled by the interface:
+   
   ```pycon
   >>> @qml.math.multi_dispatch(argnum=[0, 1])
   ... def some_function(tensor1, tensor2, option, like):
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 5616bdbe8..6b4708e13 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -599,7 +599,7 @@ class DefaultQubit(QubitDevice):
 
         # Return the full density matrix by using numpy tensor product
         if wires == self.wires:
-            density_matrix = self._tensordot(state, self._conj(state), 0)
+            density_matrix = self._tensordot(state, self._conj(state), axes=0)
             density_matrix = self._reshape(density_matrix, (2 ** len(wires), 2 ** len(wires)))
             return density_matrix
 
diff --git a/pennylane/math/is_independent.py b/pennylane/math/is_independent.py
index 5a2bea1e6..302b5bb77 100644
--- a/pennylane/math/is_independent.py
+++ b/pennylane/math/is_independent.py
@@ -203,7 +203,7 @@ def _get_random_args(args, interface, num, seed, bounds):
         if interface == "autograd":
 
             # Mark the arguments as trainable with Autograd
-            rnd_args = pnp.array(rnd_args, requires_grad=True)
+            rnd_args = [tuple(pnp.array(a, requires_grad=True) for a in arg) for arg in rnd_args]
 
     return rnd_args
 
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index a61288dbd..4b8a2d220 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -102,8 +102,9 @@ def multi_dispatch(argnum=None, tensor_list=None):
             to dispatch (i.e., the arguments that are tensors handled by an interface).
             If ``None``, dispatch over all arguments.
         tensor_lists (list[int]): a list of integers indicating which indices
-            in ``argnum`` are expected to be lists of tensors.
-            If ``None``, this option is ignored.
+            in ``argnum`` are expected to be lists of tensors. If an argument
+            marked as tensor list is not a ``tuple`` or ``list``, it is treated
+            as if it was not marked as tensor list. If ``None``, this option is ignored.
 
     Returns:
         func: A wrapped version of the function, which will automatically attempt
@@ -163,7 +164,9 @@ def multi_dispatch(argnum=None, tensor_list=None):
             dispatch_args = []
 
             for a in argnums:
-                if a in tensor_lists:
+                # Only use extend if the marked argument really
+                # is a (native) python Sequence
+                if a in tensor_lists and isinstance(args[a], (list, tuple)):
                     dispatch_args.extend(args[a])
                 else:
                     dispatch_args.append(args[a])
@@ -179,7 +182,8 @@ def multi_dispatch(argnum=None, tensor_list=None):
     return decorator
 
 
-def block_diag(values):
+@multi_dispatch(argnum=[0], tensor_list=[0])
+def block_diag(values, like=None):
     """Combine a sequence of 2D tensors to form a block diagonal tensor.
 
     Args:
@@ -203,12 +207,12 @@ def block_diag(values):
             [ 0,  0, -1, -6, -3,  0],
             [ 0,  0,  0,  0,  0,  5]])
     """
-    interface = _multi_dispatch(values)
-    values = np.coerce(values, like=interface)
-    return np.block_diag(values, like=interface)
+    values = np.coerce(values, like=like)
+    return np.block_diag(values, like=like)
 
 
-def concatenate(values, axis=0):
+@multi_dispatch(argnum=[0], tensor_list=[0])
+def concatenate(values, axis=0, like=None):
     """Concatenate a sequence of tensors along the specified axis.
 
     .. warning::
@@ -235,9 +239,7 @@ def concatenate(values, axis=0):
     <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
     array([6.00e-01, 1.00e-01, 6.00e-01, 1.00e-01, 2.00e-01, 3.00e-01, 5.00e+00, 8.00e+00, 1.01e+02], dtype=float32)>
     """
-    interface = _multi_dispatch(values)
-
-    if interface == "torch":
+    if like == "torch":
         import torch
 
         if axis is None:
@@ -248,16 +250,17 @@ def concatenate(values, axis=0):
         else:
             values = [torch.as_tensor(t) for t in values]
 
-    if interface == "tensorflow" and axis is None:
+    if like == "tensorflow" and axis is None:
         # flatten and then concatenate zero'th dimension
         # to reproduce numpy's behaviour
         values = [np.flatten(np.array(t)) for t in values]
         axis = 0
 
-    return np.concatenate(values, axis=axis, like=interface)
+    return np.concatenate(values, axis=axis, like=like)
 
 
-def diag(values, k=0):
+@multi_dispatch(argnum=[0], tensor_list=[0])
+def diag(values, k=0, like=None):
     """Construct a diagonal tensor from a list of scalars.
 
     Args:
@@ -291,15 +294,14 @@ def diag(values, k=0):
             [0.0000, 0.0000, 0.2000],
             [0.0000, 0.0000, 0.0000]])
     """
-    interface = _multi_dispatch(values)
-
     if isinstance(values, (list, tuple)):
-        values = np.stack(np.coerce(values, like=interface), like=interface)
+        values = np.stack(np.coerce(values, like=like), like=like)
 
-    return np.diag(values, k=k, like=interface)
+    return np.diag(values, k=k, like=like)
 
 
-def dot(tensor1, tensor2):
+@multi_dispatch(argnum=[0, 1])
+def dot(tensor1, tensor2, like=None):
     """Returns the matrix or dot product of two tensors.
 
     * If both tensors are 0-dimensional, elementwise multiplication
@@ -323,34 +325,34 @@ def dot(tensor1, tensor2):
     Returns:
         tensor_like: the matrix or dot product of two tensors
     """
-    interface = _multi_dispatch([tensor1, tensor2])
-    x, y = np.coerce([tensor1, tensor2], like=interface)
+    x, y = np.coerce([tensor1, tensor2], like=like)
 
-    if interface == "torch":
+    if like == "torch":
         if x.ndim == 0 and y.ndim == 0:
             return x * y
 
         if x.ndim <= 2 and y.ndim <= 2:
             return x @ y
 
-        return np.tensordot(x, y, axes=[[-1], [-2]], like=interface)
+        return np.tensordot(x, y, axes=[[-1], [-2]], like=like)
 
-    if interface == "tensorflow":
+    if like == "tensorflow":
         if len(np.shape(x)) == 0 and len(np.shape(y)) == 0:
             return x * y
 
         if len(np.shape(y)) == 1:
-            return np.tensordot(x, y, axes=[[-1], [0]], like=interface)
+            return np.tensordot(x, y, axes=[[-1], [0]], like=like)
 
         if len(np.shape(x)) == 2 and len(np.shape(y)) == 2:
             return x @ y
 
-        return np.tensordot(x, y, axes=[[-1], [-2]], like=interface)
+        return np.tensordot(x, y, axes=[[-1], [-2]], like=like)
 
-    return np.dot(x, y, like=interface)
+    return np.dot(x, y, like=like)
 
 
-def tensordot(tensor1, tensor2, axes=None):
+@multi_dispatch(argnum=[0, 1])
+def tensordot(tensor1, tensor2, axes=None, like=None):
     """Returns the tensor product of two tensors.
     In general ``axes`` specifies either the set of axes for both
     tensors that are contracted (with the first/second entry of ``axes``
@@ -376,11 +378,12 @@ def tensordot(tensor1, tensor2, axes=None):
     Returns:
         tensor_like: the tensor product of the two input tensors
     """
-    interface = _multi_dispatch([tensor1, tensor2])
-    return np.tensordot(tensor1, tensor2, axes=axes, like=interface)
+    tensor1, tensor2 = np.coerce([tensor1, tensor2], like=like)
+    return np.tensordot(tensor1, tensor2, axes=axes, like=like)
 
 
-def get_trainable_indices(values):
+@multi_dispatch(argnum=[0], tensor_list=[0])
+def get_trainable_indices(values, like=None):
     """Returns a set containing the trainable indices of a sequence of
     values.
 
@@ -403,10 +406,9 @@ def get_trainable_indices(values):
     tensor(0.0899685, requires_grad=True)
     """
     trainable = requires_grad
-    interface = _multi_dispatch(values)
     trainable_params = set()
 
-    if interface == "jax":
+    if like == "jax":
         import jax
 
         if not any(isinstance(v, jax.core.Tracer) for v in values):
@@ -420,7 +422,7 @@ def get_trainable_indices(values):
             trainable = requires_grad
 
     for idx, p in enumerate(values):
-        if trainable(p, interface=interface):
+        if trainable(p, interface=like):
             trainable_params.add(idx)
 
     return trainable_params
@@ -474,8 +476,7 @@ def safe_squeeze(tensor, axis=None, exclude_axis=None):
         or not excluded and that have size 1. If no axes are specified or excluded,
         all axes are attempted to be squeezed.
     """
-    interface = _multi_dispatch([tensor])
-    if interface == "tensorflow":
+    if get_interface(tensor) == "tensorflow":
         from tensorflow.python.framework.errors_impl import InvalidArgumentError
 
         exception = InvalidArgumentError
@@ -508,7 +509,8 @@ def safe_squeeze(tensor, axis=None, exclude_axis=None):
     return tensor
 
 
-def stack(values, axis=0):
+@multi_dispatch(argnum=[0], tensor_list=[0])
+def stack(values, axis=0, like=None):
     """Stack a sequence of tensors along the specified axis.
 
     .. warning::
@@ -537,9 +539,8 @@ def stack(values, axis=0):
            [1.00e-01, 2.00e-01, 3.00e-01],
            [5.00e+00, 8.00e+00, 1.01e+02]], dtype=float32)>
     """
-    interface = _multi_dispatch(values)
-    values = np.coerce(values, like=interface)
-    return np.stack(values, axis=axis, like=interface)
+    values = np.coerce(values, like=like)
+    return np.stack(values, axis=axis, like=like)
 
 
 def where(condition, x=None, y=None):
@@ -612,7 +613,8 @@ def where(condition, x=None, y=None):
     return np.where(condition, x, y, like=_multi_dispatch([condition, x, y]))
 
 
-def frobenius_inner_product(A, B, normalize=False):
+@multi_dispatch(argnum=[0, 1])
+def frobenius_inner_product(A, B, normalize=False, like=None):
     r"""Frobenius inner product between two matrices.
 
     .. math::
@@ -637,8 +639,7 @@ def frobenius_inner_product(A, B, normalize=False):
     >>> qml.math.frobenius_inner_product(A, B)
     3.091948202943376
     """
-    interface = _multi_dispatch([A, B])
-    A, B = np.coerce([A, B], like=interface)
+    A, B = np.coerce([A, B], like=like)
 
     inner_product = np.sum(A * B)
 
@@ -649,6 +650,7 @@ def frobenius_inner_product(A, B, normalize=False):
     return inner_product
 
 
+@multi_dispatch(argnum=[0, 2])
 def scatter_element_add(tensor, index, value, like=None):
     """In-place addition of a multidimensional value over various
     indices of a tensor.
@@ -682,8 +684,7 @@ def scatter_element_add(tensor, index, value, like=None):
     if len(np.shape(tensor)) == 0 and index == ():
         return tensor + value
 
-    interface = like or _multi_dispatch([tensor, value])
-    return np.scatter_element_add(tensor, index, value, like=interface)
+    return np.scatter_element_add(tensor, index, value, like=like)
 
 
 def unwrap(values, max_depth=None):


Test patch:
diff --git a/tests/math/test_basic_math.py b/tests/math/test_basic_math.py
index 9a098eb9f..a592749a2 100644
--- a/tests/math/test_basic_math.py
+++ b/tests/math/test_basic_math.py
@@ -115,7 +115,7 @@ class TestFrobeniusInnerProduct:
 
     def test_frobenius_inner_product_gradient(self):
         """Test that the calculated gradient is correct."""
-        A = np.array([[1.0, 2.3], [-1.3, 2.4]])
+        A = onp.array([[1.0, 2.3], [-1.3, 2.4]])
         B = torch.autograd.Variable(torch.randn(2, 2).type(torch.float), requires_grad=True)
         result = fn.frobenius_inner_product(A, B)
         result.backward()
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index 2be238b7f..8434d7b73 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Unit tests for the TensorBox functional API in pennylane.fn.fn
 """
+from functools import partial
 import itertools
 import numpy as onp
 import pytest
@@ -322,7 +323,8 @@ class TestConcatenate:
     """Tests for the concatenate function"""
 
     def test_concatenate_array(self):
-        """Test that concatenate, called without the axis arguments, concatenates across the 0th dimension"""
+        """Test that concatenate, called without the axis arguments,
+        concatenates across the 0th dimension"""
         t1 = [0.6, 0.1, 0.6]
         t2 = np.array([0.1, 0.2, 0.3])
         t3 = onp.array([5.0, 8.0, 101.0])
@@ -332,7 +334,8 @@ class TestConcatenate:
         assert np.all(res == np.concatenate([t1, t2, t3]))
 
     def test_concatenate_jax(self):
-        """Test that concatenate, called without the axis arguments, concatenates across the 0th dimension"""
+        """Test that concatenate, called without the axis arguments,
+        concatenates across the 0th dimension"""
         t1 = jnp.array([5.0, 8.0, 101.0])
         t2 = jnp.array([0.6, 0.1, 0.6])
         t3 = jnp.array([0.1, 0.2, 0.3])
@@ -340,8 +343,9 @@ class TestConcatenate:
         res = fn.concatenate([t1, t2, t3])
         assert jnp.all(res == jnp.concatenate([t1, t2, t3]))
 
-    def test_stack_tensorflow(self):
-        """Test that concatenate, called without the axis arguments, concatenates across the 0th dimension"""
+    def test_concatenate_tensorflow(self):
+        """Test that concatenate, called without the axis arguments,
+        concatenates across the 0th dimension"""
         t1 = tf.constant([0.6, 0.1, 0.6])
         t2 = tf.Variable([0.1, 0.2, 0.3])
         t3 = onp.array([5.0, 8.0, 101.0])
@@ -350,8 +354,9 @@ class TestConcatenate:
         assert isinstance(res, tf.Tensor)
         assert np.all(res.numpy() == np.concatenate([t1.numpy(), t2.numpy(), t3]))
 
-    def test_stack_torch(self):
-        """Test that concatenate, called without the axis arguments, concatenates across the 0th dimension"""
+    def test_concatenate_torch(self):
+        """Test that concatenate, called without the axis arguments,
+        concatenates across the 0th dimension"""
         t1 = onp.array([5.0, 8.0, 101.0], dtype=np.float64)
         t2 = torch.tensor([0.6, 0.1, 0.6], dtype=torch.float64)
         t3 = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float64)
@@ -363,7 +368,7 @@ class TestConcatenate:
     @pytest.mark.parametrize(
         "t1", [onp.array([[1], [2]]), torch.tensor([[1], [2]]), tf.constant([[1], [2]])]
     )
-    def test_stack_axis(self, t1):
+    def test_concatenate_axis(self, t1):
         """Test that passing the axis argument allows for concatenating along
         a different axis"""
         t2 = onp.array([[3], [4]])
@@ -546,7 +551,9 @@ class TestDot:
 
 
 class TestTensordotTorch:
-    """Tests for the tensor product function in torch."""
+    """Tests for the tensor product function in torch.
+    This test is required because the functionality of tensordot for Torch
+    is being patched in PennyLane, as compared to autoray."""
 
     v1 = torch.tensor([0.1, 0.5, -0.9, 1.0, -4.2, 0.1], dtype=torch.float64)
     v2 = torch.tensor([4.3, -1.2, 8.2, 0.6, -4.2, -11.0], dtype=torch.float64)
@@ -762,6 +769,89 @@ class TestTensordotTorch:
         assert fn.allclose(fn.tensordot(self.T1, M, axes=[axes1, axes2]), expected)
 
 
+class TestTensordotDifferentiability:
+
+    v0 = np.array([0.1, 5.3, -0.9, 1.1])
+    v1 = np.array([0.5, -1.7, -2.9, 0.0])
+    v2 = np.array([-0.4, 9.1, 1.6])
+    exp_shapes = ((len(v0), len(v2), len(v0)), (len(v0), len(v2), len(v2)))
+    exp_jacs = (np.zeros(exp_shapes[0]), np.zeros(exp_shapes[1]))
+    for i in range(len(v0)):
+        exp_jacs[0][i, :, i] = v2
+    for i in range(len(v2)):
+        exp_jacs[1][:, i, i] = v0
+
+    def test_autograd(self):
+        """Tests differentiability of tensordot with Autograd."""
+        v0 = np.array(self.v0, requires_grad=True)
+        v1 = np.array(self.v1, requires_grad=True)
+        v2 = np.array(self.v2, requires_grad=True)
+
+        # Test inner product
+        jac = qml.jacobian(partial(fn.tensordot, axes=[0, 0]), argnum=(0, 1))(v0, v1)
+        assert all(fn.allclose(jac[i], _v) for i, _v in enumerate([v1, v0]))
+
+        # Test outer product
+        jac = qml.jacobian(partial(fn.tensordot, axes=0), argnum=(0, 1))(v0, v2)
+        assert all(fn.shape(jac[i]) == self.exp_shapes[i] for i in [0, 1])
+        assert all(fn.allclose(jac[i], self.exp_jacs[i]) for i in [0, 1])
+
+    def test_torch(self):
+        """Tests differentiability of tensordot with Torch."""
+        jac_fn = torch.autograd.functional.jacobian
+
+        v0 = torch.tensor(self.v0, requires_grad=True, dtype=torch.float64)
+        v1 = torch.tensor(self.v1, requires_grad=True, dtype=torch.float64)
+        v2 = torch.tensor(self.v2, requires_grad=True, dtype=torch.float64)
+
+        # Test inner product
+        jac = jac_fn(partial(fn.tensordot, axes=[[0], [0]]), (v0, v1))
+        assert all(fn.allclose(jac[i], _v) for i, _v in enumerate([v1, v0]))
+
+        # Test outer product
+        jac = jac_fn(partial(fn.tensordot, axes=0), (v0, v2))
+        assert all(fn.shape(jac[i]) == self.exp_shapes[i] for i in [0, 1])
+        assert all(fn.allclose(jac[i], self.exp_jacs[i]) for i in [0, 1])
+
+    def test_jax(self):
+        """Tests differentiability of tensordot with JAX."""
+        jac_fn = jax.jacobian
+
+        v0 = jnp.array(self.v0)
+        v1 = jnp.array(self.v1)
+        v2 = jnp.array(self.v2)
+
+        # Test inner product
+        jac = jac_fn(partial(fn.tensordot, axes=[[0], [0]]), argnums=(0, 1))(v0, v1)
+        assert all(fn.allclose(jac[i], _v) for i, _v in enumerate([v1, v0]))
+
+        # Test outer product
+        jac = jac_fn(partial(fn.tensordot, axes=0), argnums=(0, 1))(v0, v2)
+        assert all(fn.shape(jac[i]) == self.exp_shapes[i] for i in [0, 1])
+        assert all(fn.allclose(jac[i], self.exp_jacs[i]) for i in [0, 1])
+
+    def test_tensorflow(self):
+        """Tests differentiability of tensordot with TensorFlow."""
+
+        def jac_fn(func, args):
+            with tf.GradientTape() as tape:
+                out = func(*args)
+            return tape.jacobian(out, args)
+
+        v0 = tf.Variable(self.v0, dtype=tf.float64)
+        v1 = tf.Variable(self.v1, dtype=tf.float64)
+        v2 = tf.Variable(self.v2, dtype=tf.float64)
+
+        # Test inner product
+        jac = jac_fn(partial(fn.tensordot, axes=[[0], [0]]), (v0, v1))
+        assert all(fn.allclose(jac[i], _v) for i, _v in enumerate([v1, v0]))
+
+        # Test outer product
+        jac = jac_fn(partial(fn.tensordot, axes=0), (v0, v2))
+        assert all(fn.shape(jac[i]) == self.exp_shapes[i] for i in [0, 1])
+        assert all(fn.allclose(jac[i], self.exp_jacs[i]) for i in [0, 1])
+
+
 # the following test data is of the form
 # [original shape, axis to expand, new shape]
 expand_dims_test_data = [
@@ -1362,14 +1452,14 @@ class TestScatterElementAdd:
         x = np.array(self.x, requires_grad=True)
         y = np.array(self.y, requires_grad=True)
 
-        def cost(weights):
+        def cost(*weights):
             return fn.scatter_element_add(weights[0], self.index, weights[1] ** 2)
 
-        res = cost([x, y])
+        res = cost(x, y)
         assert isinstance(res, np.ndarray)
         assert fn.allclose(res, self.expected_val)
 
-        grad = qml.grad(lambda weights: cost(weights)[self.index[0], self.index[1]])([x, y])
+        grad = qml.grad(lambda *weights: cost(*weights)[self.index[0], self.index[1]])(x, y)
         assert fn.allclose(grad[0], self.expected_grad_x)
         assert fn.allclose(grad[1], self.expected_grad_y)
 
@@ -1492,20 +1582,20 @@ class TestScatterElementAddMultiValue:
         x = np.array(self.x, requires_grad=True)
         y = np.array(self.y, requires_grad=True)
 
-        def cost(weights):
+        def cost(*weights):
             return fn.scatter_element_add(
                 weights[0], self.indices, [fn.sin(weights[1] / 2), weights[1] ** 2]
             )
 
-        res = cost([x, y])
+        res = cost(x, y)
         assert isinstance(res, np.ndarray)
         assert fn.allclose(res, self.expected_val)
 
         scalar_cost = (
-            lambda weights: cost(weights)[self.indices[0][0], self.indices[1][0]]
-            + cost(weights)[self.indices[0][1], self.indices[1][1]]
+            lambda *weights: cost(*weights)[self.indices[0][0], self.indices[1][0]]
+            + cost(*weights)[self.indices[0][1], self.indices[1][1]]
         )
-        grad = qml.grad(scalar_cost)([x, y])
+        grad = qml.grad(scalar_cost)(x, y)
         assert fn.allclose(grad[0], self.expected_grad_x)
         assert fn.allclose(grad[1], self.expected_grad_y)
 
@@ -1871,7 +1961,7 @@ class TestBlockDiagDiffability:
             np.array([[x, 1.2 * y], [x ** 2 - y / 3, -x / y]]),
         ]
         f = lambda x, y: fn.block_diag(tensors(x, y))
-        x, y = 0.2, 1.5
+        x, y = np.array([0.2, 1.5], requires_grad=True)
         res = qml.jacobian(f)(x, y)
         exp = self.expected(x, y)
         assert fn.allclose(res[0], exp[0])
@@ -2029,12 +2119,16 @@ class TestUnwrap:
             unwrapped_params = qml.math.unwrap(params)
             return np.sum(np.sin(params[0] * params[2])) + params[1]
 
-        values = [onp.array([0.1, 0.2]), np.tensor(0.1, dtype=np.float64), np.tensor([0.5, 0.2])]
-        grad = qml.grad(cost_fn)(*values)
+        values = [
+            onp.array([0.1, 0.2]),
+            np.tensor(0.1, dtype=np.float64, requires_grad=True),
+            np.tensor([0.5, 0.2], requires_grad=True),
+        ]
+        grad = qml.grad(cost_fn, argnum=[1, 2])(*values)
 
         expected = [np.array([0.1, 0.2]), 0.1, np.array([0.5, 0.2])]
         assert all(np.allclose(a, b) for a, b in zip(unwrapped_params, expected))
-        assert all(not isinstance(a, ArrayBox) for a in unwrapped_params)
+        assert not any(isinstance(a, ArrayBox) for a in unwrapped_params)
 
     def test_autograd_unwrapping_backward_nested(self):
         """Test that a sequence of Autograd values is properly unwrapped
diff --git a/tests/math/test_is_independent.py b/tests/math/test_is_independent.py
index 6953b584d..cc2af6e6a 100644
--- a/tests/math/test_is_independent.py
+++ b/tests/math/test_is_independent.py
@@ -56,8 +56,8 @@ dependent_lambdas = [
     lambda x: x if abs(x) < 1e-5 else 0.0,  # x*delta for x=0 is okay
     lambda x: 1.0 if x > 0 else 0.0,  # Heaviside is okay numerically
     lambda x: 1.0 if x > 0 else 0.0,  # Heaviside is okay numerically
-    lambda x: qml.math.log(1 + qml.math.exp(1000.0 * x)) / 1000.0,  # Softplus is okay
-    lambda x: qml.math.log(1 + qml.math.exp(1000.0 * x)) / 1000.0,  # Softplus is okay
+    lambda x: qml.math.log(1 + qml.math.exp(100.0 * x)) / 100.0,  # Softplus is okay
+    lambda x: qml.math.log(1 + qml.math.exp(100.0 * x)) / 100.0,  # Softplus is okay
 ]
 
 args_dependent_lambdas = [
@@ -167,7 +167,7 @@ class TestIsIndependentAutograd:
         dependent_circuit,
         np.array,
         lambda x: np.array(x * 0.0),
-        lambda x: (1 + qml.math.tanh(1000 * x)) / 2,
+        lambda x: (1 + qml.math.tanh(100 * x)) / 2,
         *dependent_lambdas,
     ]
 
@@ -519,7 +519,7 @@ if have_torch:
 
         dependent_functions = [
             dependent_circuit,
-            torch.tensor,
+            torch.as_tensor,
             lambda x: (1 + qml.math.tanh(1000 * x)) / 2,
             *dependent_lambdas,
         ]