Repository: PennyLaneAI/pennylane

Issue:
[BUG] Split non commuting fails with lightning.kokkos
### Expected behavior

Split non commuting works for lightning kokkos

### Actual behavior

An exception is thrown.

### Additional information

Not an urgent bug, just looking at some tests I have for catalyst transforms.

### Source code

```shell
import jax
import pennylane as qml
import pennylane.transforms

def test_split_non_commuting():
    @qml.transforms.split_non_commuting
    @qml.qnode(qml.device("lightning.kokkos", wires=6), interface="jax")
    def interpreted():
        qml.Hadamard(1)
        qml.Hadamard(0)
        qml.PauliZ(0)
        qml.Hadamard(3)
        qml.Hadamard(5)
        qml.T(5)
        return (
            qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
            qml.expval(qml.PauliX(0)),
            qml.expval(qml.PauliZ(1)),
            qml.expval(qml.PauliX(1) @ qml.PauliX(4)),
            qml.expval(qml.PauliX(3)),
            qml.expval(qml.PauliY(5)),
        )

    interpreted()

test_split_non_commuting()
```


### Tracebacks

```shell
python bug.py 
Kokkos::OpenMP::initialize WARNING: OMP_PROC_BIND environment variable not set
  In general, for best performance with OpenMP 4.0 or better set OMP_PROC_BIND=spread and OMP_PLACES=threads
  For best performance with OpenMP 3.1 set OMP_PROC_BIND=true
  For unit testing set OMP_PROC_BIND=false

Traceback (most recent call last):
  File "/home/erick.ochoalopez/Code/cataliist/bug.py", line 26, in <module>
    test_split_non_commuting()
  File "/home/erick.ochoalopez/Code/cataliist/bug.py", line 8, in test_split_non_commuting
    def interpreted():
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/pennylane/transforms/core/transform_dispatcher.py", line 96, in __call__
    return self._qnode_transform(obj, targs, tkwargs)
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/pennylane/transforms/core/transform_dispatcher.py", line 208, in default_qnode_transform
    qnode = copy.deepcopy(qnode)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/erick.ochoalopez/.pyenv/versions/3.10.11/lib/python3.10/copy.py", line 161, in deepcopy
    rv = reductor(4)
TypeError: cannot pickle 'pennylane_lightning_kokkos.lightning_kokkos_qubit_ops.LightningKokkos_C128' object
```


### System information

```shell
Python 3.10.11 (main, May 13 2023, 12:07:51) [GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import pennylane as qml
>>> qml.about()
Name: PennyLane
Version: 0.34.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/erick.ochoalopez/Code/env/lib/python3.10/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: pennylane-catalyst, PennyLane-Lightning, PennyLane-Lightning-Kokkos

Platform info:           Linux-5.15.0-84-generic-x86_64-with-glibc2.31
Python version:          3.10.11
Numpy version:           1.26.1
Scipy version:           1.11.3
Installed devices:
- default.gaussian (PennyLane-0.34.0.dev0)
- default.mixed (PennyLane-0.34.0.dev0)
- default.qubit (PennyLane-0.34.0.dev0)
- default.qubit.autograd (PennyLane-0.34.0.dev0)
- default.qubit.jax (PennyLane-0.34.0.dev0)
- default.qubit.legacy (PennyLane-0.34.0.dev0)
- default.qubit.tf (PennyLane-0.34.0.dev0)
- default.qubit.torch (PennyLane-0.34.0.dev0)
- default.qutrit (PennyLane-0.34.0.dev0)
- null.qubit (PennyLane-0.34.0.dev0)
- lightning.kokkos (PennyLane-Lightning-Kokkos-0.32.0)
- lightning.qubit (PennyLane-Lightning-0.32.0)
>>>
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/gradients/vjp.py b/pennylane/gradients/vjp.py
--- a/pennylane/gradients/vjp.py
+++ b/pennylane/gradients/vjp.py
@@ -47,7 +47,7 @@ def _all_close_to_zero(dy):
         return qml.math.allclose(dy, 0)
 
     # call this method recursively
-    return qml.math.all(qml.math.stack([_all_close_to_zero(dy_) for dy_ in dy]))
+    return all(_all_close_to_zero(dy_) for dy_ in dy)
 
 
 def compute_vjp_single(dy, jac, num=None):
@@ -75,7 +75,7 @@ def compute_vjp_single(dy, jac, num=None):
         >>> compute_vjp_single(dy, jac)
         np.array([0.2])
 
-    2. For a single parameter and a single measurment with shape (e.g. probs):
+    2. For a single parameter and a single measurement with shape (e.g. probs):
 
     .. code-block:: pycon
 
@@ -115,16 +115,8 @@ def compute_vjp_single(dy, jac, num=None):
     if not isinstance(dy_row, np.ndarray):
         jac = _convert(jac, dy_row)
 
-    try:
-        if _all_close_to_zero(dy):
-            # If the dy vector is zero, then the
-            # corresponding element of the VJP will be zero.
-            num_params = len(jac) if isinstance(jac, tuple) else 1
-
-            res = qml.math.convert_like(np.zeros(num_params), dy)
-            return qml.math.cast_like(res, dy)
-    except (AttributeError, TypeError):
-        pass
+    # Note: For generality, all exception type warnings are disabled.
+    # TODO: Excplictly catalogue and update raises for known types.
 
     # Single measurement with a single param
     if not isinstance(jac, (tuple, autograd.builtins.SequenceBox)):
@@ -136,7 +128,12 @@ def compute_vjp_single(dy, jac, num=None):
         if num == 1:
             jac = qml.math.squeeze(jac)
         jac = qml.math.reshape(jac, (-1, 1))
-        res = qml.math.tensordot(jac, dy_row, [[0], [0]])
+        try:
+            res = dy_row @ jac
+
+        except Exception:  # pylint: disable=broad-except
+            res = qml.math.tensordot(jac, dy_row, [[0], [0]])
+
     # Single measurement with multiple params
     else:
         # No trainable parameters (adjoint)
@@ -146,12 +143,19 @@ def compute_vjp_single(dy, jac, num=None):
         # Single measurement with no dimension e.g. expval
         if num == 1:
             jac = qml.math.reshape(qml.math.stack(jac), (1, -1))
-            res = qml.math.tensordot(jac, dy_row, [[0], [0]])
+            try:
+                res = dy_row @ jac
+            except Exception:  # pylint: disable=broad-except
+                res = qml.math.tensordot(jac, dy_row, [[0], [0]])
 
         # Single measurement with dimension e.g. probs
         else:
             jac = qml.math.stack(jac)
-            res = qml.math.tensordot(jac, dy_row, [[1], [0]])
+            try:
+                res = jac @ dy_row
+            except Exception:  # pylint: disable=broad-except
+                res = qml.math.tensordot(jac, dy_row, [[1], [0]])
+
     return res
 
 
@@ -202,14 +206,34 @@ def compute_vjp_multi(dy, jac, num=None):
         res = qml.math.sum(qml.math.stack(res), axis=0)
     # Multiple parameters
     else:
-        res = []
-        for d, j_ in zip(dy, jac):
-            sub_res = []
-            for j in j_:
-                sub_res.append(qml.math.squeeze(compute_vjp_single(d, j, num=num)))
-            res.append(sub_res)
-        res = qml.math.stack([qml.math.stack(r) for r in res])
-        res = qml.math.sum(res, axis=0)
+        try:
+            dy_interface = qml.math.get_interface(dy[0])
+            # dy  -> (i,j)      observables, entries per observable
+            # jac -> (i,k,j)    observables, parameters, entries per observable
+            # Contractions over observables and entries per observable
+            dy_shape = qml.math.shape(dy)
+            if len(dy_shape) > 1:  # multiple values exist per observable output
+                return qml.math.array(qml.math.einsum("ij,i...j", dy, jac), like=dy[0])
+
+            if dy_interface == "tensorflow":
+                # TF needs a different path for Hessian support
+                return qml.math.array(
+                    qml.math.einsum("i,i...", dy, jac, like=dy[0]), like=dy[0]
+                )  # Scalar value per observable output
+            return qml.math.array(
+                qml.math.einsum("i,i...", dy, jac), like=dy[0]
+            )  # Scalar value per observable output
+        # NOTE: We want any potential failure to fall back here, so catch every exception type
+        # TODO: Catalogue and update for expected exception types
+        except Exception:  # pylint: disable=broad-except
+            res = []
+            for d, j_ in zip(dy, jac):
+                sub_res = []
+                for j in j_:
+                    sub_res.append(qml.math.squeeze(compute_vjp_single(d, j, num=num)))
+                res.append(sub_res)
+            res = qml.math.stack([qml.math.stack(r) for r in res])
+            res = qml.math.sum(res, axis=0)
     return res
 
 
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -540,6 +540,13 @@ def einsum(indices, *operands, like=None, optimize=None):
     if optimize is None or like == "torch":
         # torch einsum doesn't support the optimize keyword argument
         return np.einsum(indices, *operands, like=like)
+    if like == "tensorflow":
+        # Unpacking and casting necessary for higher order derivatives,
+        # and avoiding implicit fp32 down-conversions.
+        op1, op2 = operands
+        op1 = array(op1, like=op1[0], dtype=op1[0].dtype)
+        op2 = array(op2, like=op2[0], dtype=op2[0].dtype)
+        return np.einsum(indices, op1, op2, like=like)
     return np.einsum(indices, *operands, like=like, optimize=optimize)
 
 


Test patch:
diff --git a/tests/gradients/core/test_vjp.py b/tests/gradients/core/test_vjp.py
--- a/tests/gradients/core/test_vjp.py
+++ b/tests/gradients/core/test_vjp.py
@@ -117,7 +117,7 @@ def test_zero_dy_single_measurement_single_params(self):
 
     def test_zero_dy_single_measurement_multi_params(self):
         """A zero dy vector will return a zero matrix"""
-        dy = np.zeros([2])
+        dy = np.zeros(1)
         jac = tuple([np.array(0.1), np.array(0.2)])
 
         vjp = qml.gradients.compute_vjp_single(dy, jac)
diff --git a/tests/gradients/finite_diff/test_finite_difference.py b/tests/gradients/finite_diff/test_finite_difference.py
--- a/tests/gradients/finite_diff/test_finite_difference.py
+++ b/tests/gradients/finite_diff/test_finite_difference.py
@@ -477,6 +477,15 @@ def __init__(self, val):
             def __mul__(self, other):
                 return SpecialObject(self.val * other)
 
+            def __rmul__(self, other):
+                return SpecialObject(other * self.val)
+
+            def __matmul__(self, other):
+                return SpecialObject(self.val @ other)
+
+            def __rmatmul__(self, other):
+                return SpecialObject(other @ self.val)
+
             def __add__(self, other):
                 new = self.val + (other.val if isinstance(other, self.__class__) else other)
                 return SpecialObject(new)