Repository: PennyLaneAI/pennylane

Issue:
[BUG] `CNOT` and `PauliZ` do not work with large batched states and tensorflow
### Expected behavior

When using tensorflow, batched data, and states with 8 or more wires, I don't expect any error to be raised.

### Actual behavior

A cryptic error is raised.

### Additional information

First reported [on the forum](https://discuss.pennylane.ai/t/abnormal-operation-of-a-specific-circuit-possibly-caused-by-tf-batch/3712)

### Source code

```shell
import pennylane as qml
from pennylane import numpy as np

import tensorflow as tf
from pennylane.templates.embeddings import AmplitudeEmbedding

dev = qml.device("default.qubit", wires=8)

@qml.qnode(dev, interface="tf")
def ancillary_qcnn_circuit(inputs):
    AmplitudeEmbedding(features=inputs, wires=range(4), normalize=True)    
    qml.CNOT(wires=[0,1])
    qml.PauliZ(1)
    qml.Toffoli(wires=[0, 2, 4])
    qml.Toffoli(wires=[0, 2, 5])
    qml.Toffoli(wires=[0, 2, 6])
    qml.Toffoli(wires=[0, 2, 7])
    return [qml.expval(qml.PauliZ(i)) for i in range(4, 8)]

params = np.random.rand(3,16)
ancillary_qcnn_circuit(tf.Variable(params))
```


### Tracebacks

```shell
File "/var/folders/w9/9mnrk11j25b3klztdzb01bqr0000gq/T/ipykernel_86620/2335111251.py", line 23, in <module>
    ancillary_qcnn_circuit(tf.Variable(params))
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/qnode.py", line 1030, in __call__
    res = qml.execute(
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/interfaces/execution.py", line 631, in execute
    results = inner_execute(tapes)
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/interfaces/execution.py", line 252, in inner_execute
    return cached_device_execution(tapes)
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/interfaces/execution.py", line 374, in wrapper
    res = list(fn(tuple(execution_tapes.values()), **kwargs))
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/devices/default_qubit.py", line 478, in execute
    results = tuple(
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/devices/default_qubit.py", line 479, in <genexpr>
    simulate(
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/devices/qubit/simulate.py", line 228, in simulate
    state, is_state_batched = get_final_state(circuit, debugger=debugger, interface=interface)
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/devices/qubit/simulate.py", line 123, in get_final_state
    state = apply_operation(op, state, is_state_batched=is_state_batched, debugger=debugger)
  File "/Users/matthews/.pyenv/versions/3.9.13/lib/python3.9/functools.py", line 888, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/devices/qubit/apply_operation.py", line 257, in apply_cnot
    return apply_operation_tensordot(op, state)
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/devices/qubit/apply_operation.py", line 133, in apply_operation_tensordot
    tdot = math.tensordot(reshaped_mat, state, axes=axes)
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/math/multi_dispatch.py", line 151, in wrapper
    return fn(*args, **kwargs)
  File "/Users/matthews/src/github.com/PennyLaneAI/pennylane/pennylane/math/multi_dispatch.py", line 389, in tensordot
    return np.tensordot(tensor1, tensor2, axes=axes, like=like)
  File "/Users/matthews/.pyenv/versions/3.9.13/envs/pl/lib/python3.9/site-packages/autoray/autoray.py", line 80, in do
    return get_lib_fn(backend, fn)(*args, **kwargs)
  File "/Users/matthews/.pyenv/versions/3.9.13/envs/pl/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/Users/matthews/.pyenv/versions/3.9.13/envs/pl/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 5883, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__MatMul_device_/job:localhost/replica:0/task:0/device:CPU:0}} Matrix size-incompatible: In[0]: [4,4], In[1]: [6,128] [Op:MatMul] name:
```


### System information

```shell
pl dev
```


### 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/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -236,7 +236,7 @@ def apply_pauliz(op: qml.PauliZ, state, is_state_batched: bool = False, debugger
     n_dim = math.ndim(state)
 
     if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-        return apply_operation_tensordot(op, state)
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     sl_0 = _get_slice(0, axis, n_dim)
     sl_1 = _get_slice(1, axis, n_dim)
@@ -254,7 +254,7 @@ def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=Non
     n_dim = math.ndim(state)
 
     if n_dim >= 9 and math.get_interface(state) == "tensorflow":
-        return apply_operation_tensordot(op, state)
+        return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
     sl_0 = _get_slice(0, control_axes, n_dim)
     sl_1 = _get_slice(1, control_axes, n_dim)


Test patch:
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -1187,15 +1187,57 @@ def test_with_torch(self, batch_dim):
 
 
 @pytest.mark.tf
-@pytest.mark.parametrize("op", (qml.PauliZ(8), qml.CNOT((5, 6))))
-def test_tf_large_state(op):
-    """Tests that custom kernels that use slicing fall back to a different method when
-    the state has a large number of wires."""
-    import tensorflow as tf
-
-    state = np.zeros([2] * 10)
-    state = tf.Variable(state)
-    new_state = apply_operation(op, state)
-
-    # still all zeros.  Mostly just making sure error not raised
-    assert qml.math.allclose(state, new_state)
+class TestLargeTFCornerCases:
+    """Test large corner cases for tensorflow."""
+
+    @pytest.mark.parametrize("op", (qml.PauliZ(8), qml.CNOT((5, 6))))
+    def test_tf_large_state(self, op):
+        """Tests that custom kernels that use slicing fall back to a different method when
+        the state has a large number of wires."""
+        import tensorflow as tf
+
+        state = np.zeros([2] * 10)
+        state = tf.Variable(state)
+        new_state = apply_operation(op, state)
+
+        # still all zeros.  Mostly just making sure error not raised
+        assert qml.math.allclose(state, new_state)
+
+    def test_cnot_large_batched_state_tf(self):
+        """Test that CNOT with large batched states works as expected."""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=8)
+
+        @qml.qnode(dev, interface="tf")
+        def ancillary_qcnn_circuit(inputs):
+            qml.AmplitudeEmbedding(features=inputs, wires=range(4), normalize=True)
+            qml.CNOT(wires=[0, 1])
+            qml.PauliZ(1)
+            qml.Toffoli(wires=[0, 2, 4])
+            qml.Toffoli(wires=[0, 2, 5])
+            qml.Toffoli(wires=[0, 2, 6])
+            qml.Toffoli(wires=[0, 2, 7])
+            return [qml.expval(qml.PauliZ(i)) for i in range(4, 8)]
+
+        batch_size = 3
+        params = np.random.rand(batch_size, 16)
+        result = ancillary_qcnn_circuit(tf.Variable(params))
+        assert qml.math.shape(result) == (4, batch_size)
+
+    def test_pauliz_large_batched_state_tf(self):
+        """Test that PauliZ with large batched states works as expected."""
+        import tensorflow as tf
+
+        @qml.qnode(qml.device("default.qubit"), interface="tf")
+        def circuit(init_state):
+            qml.StatePrep(init_state, wires=range(8))
+            qml.PauliX(0)
+            qml.PauliZ(0)
+            return qml.state()
+
+        states = np.zeros((3, 256))
+        states[:, 0] = 1.0
+        results = circuit(tf.Variable(states))
+        assert qml.math.shape(results) == (3, 256)
+        assert np.array_equal(results[:, 128], [-1.0 + 0.0j] * 3)