fix devices.qubit.apply_operation with large tf states

[albi3ro]

Tensorflow doesn't support slicing on states with more than 8 dimensions. So if we try and run the code:

state = create_initial_state(range(10), like='tensorflow')
apply_operation(qml.PauliZ(9), state)

We get:

---------------------------------------------------------------------------
UnimplementedError                        Traceback (most recent call last)
Cell In[7], line 2
      1 state = create_initial_state(range(10), like='tensorflow')
----> 2 apply_operation(qml.PauliZ(9), state)

File Python.framework/Versions/3.10/lib/python3.10/functools.py:889, in singledispatch.<locals>.wrapper(*args, **kw)
    885 if not args:
    886     raise TypeError(f'{funcname} requires at least '
    887                     '1 positional argument')
--> 889 return dispatch(args[0].__class__)(*args, **kw)

File pennylane/pennylane/devices/qubit/apply_operation.py:183, in apply_pauliz(op, state)
    180 sl_0 = _get_slice(0, op.wires[0], n_wires)
    181 sl_1 = _get_slice(1, op.wires[0], n_wires)
--> 183 state1 = math.multiply(state[sl_1], -1)
    184 return math.stack([state[sl_0], state1], axis=op.wires[0])

File /lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File lib/python3.10/site-packages/tensorflow/python/framework/ops.py:7215, in raise_from_not_ok_status(e, name)
   7213 def raise_from_not_ok_status(e, name):
   7214   e.message += (" name: " + name if name is not None else "")
-> 7215   raise core._status_to_exception(e) from None

UnimplementedError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:CPU:0}} Unhandled input dimensions 10 [Op:StridedSlice] name: strided_slice

Since this case is exceptional, and we do not want to slow down the standard logical slow, I implemented fallback logic using try-except blocks.

If the error is a tf.errors.UnimplementedError, we fall back to apply_operation_tensordot.

This will probably have performance implications for applying PauliZ and CNOT on large states, but at least it won't affect the performance in cases where the error does not occur.

Blocks testing for #3862

[timmysilv]

apply_times.txt
imo, this can just be a conditional check right at the start of the function:

n_dim = qml.math.ndim(state)
if n_dim >= 9 and qml.math.get_interface(state) == "tensorflow":
    return apply_operation_tensordot(op, state)

I think this is better for a few reasons:

• For smaller states, this will evaluate to false immediately without dispatching to math. For larger states, this won't be the bottleneck anyway
• The code is much clearer
• This only checks for the exact scenario that you know causes the bug. If something we didn't foresee happens, we might silently default to tensordot without realizing or knowing why
• I suspected that try-except is slow/not preferred in general, so I benchmarked the difference between try-except and using get_interface (with the lazy evaluation suggested above) over a range of system sizes. I attached a text file with results, but tl;dr the time is almost identical for low systems (try-except might be the teeny tiniest bit faster below 9 qubits because of the get_interface call, but it's a matter of microseconds), but notably faster for 9+ qubit systems, around 100 microseconds faster.

[albi3ro]

@timmysilv Thanks for running those benchmarks. I've updated the code to to check first.

[codecov]

Codecov Report

Merging #3884 (4106c59) into master (bdbd929) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #3884   +/-   ##
=======================================
  Coverage   99.71%   99.71%           
=======================================
  Files         341      341           
  Lines       29810    29814    +4     
=======================================
+ Hits        29724    29728    +4     
  Misses         86       86           

Impacted Files	Coverage Δ
pennylane/devices/qubit/apply_operation.py	100.00% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[mudit2812]

Looks good Christina! Are there any other interfaces such as torch that might also share the slicing trouble?

[albi3ro]

[sc-34873]

[albi3ro]

Looks good Christina! Are there any other interfaces such as torch that might also share the slicing trouble?

This problem was found when writing tests for #3862 . All other interfaces passed those tests.

[mudit2812]

Thanks Christina!

[timmysilv]

i realize this isn't the most fitting changelog entry anymore... maybe let's sneak a touch-up into another PR?