fix devices.qubit.apply_operation with large tf states

doc/releases/changelog-dev.md

@@ -111,7 +111,10 @@
 
 * Made `qml.OrbitalRotation` and consequently `qml.GateFabric` consistent with the interleaved Jordan-Wigner ordering.
   [(#3861)](https://github.com/PennyLaneAI/pennylane/pull/3861)
-
+
+* `qml.devices.qubit.apply_operation` catches the `tf.errors.UnimplementedError` that occurs when `PauliZ` or `CNOT` gates
+  are applied to a large (>8 wires) tensorflow state. When that occurs, the logic falls back to the tensordot logic instead.
+  [(#3884)](https://github.com/PennyLaneAI/pennylane/pull/3884/)
 
 <h3>Contributors</h3>
 

pennylane/devices/qubit/apply_operation.py

@@ -201,12 +201,17 @@ def apply_paulix(op: qml.PauliX, state, is_state_batched: bool = False):
 @apply_operation.register
 def apply_pauliz(op: qml.PauliZ, state, is_state_batched: bool = False):
     """Apply pauliz to state."""
+
     axis = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state)
+
     sl_0 = _get_slice(0, axis, n_dim)
     sl_1 = _get_slice(1, axis, n_dim)
 
+    # must be first state and then -1 because it breaks otherwise
     state1 = math.multiply(state[sl_1], -1)
     return math.stack([state[sl_0], state1], axis=axis)
 
@@ -218,6 +223,9 @@ def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False):
     control_axes = op.wires[0] + is_state_batched
     n_dim = math.ndim(state)
 
+    if n_dim >= 9 and math.get_interface(state) == "tensorflow":
+        return apply_operation_tensordot(op, state)
+
     sl_0 = _get_slice(0, control_axes, n_dim)
     sl_1 = _get_slice(1, control_axes, n_dim)
 

tests/devices/qubit/test_apply_operation.py

@@ -453,3 +453,18 @@ def test_double_excitation(self, method):
             state_v2 = method(d_op, state_v2)
 
         assert qml.math.allclose(state_v1, state_v2)
+
+
+@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)