qc_qmc fails with pickling error on lightning.qubit

[licedric]

Describe the bug
The qc_qmc notebook fails if run as-is with the following stack trace:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File <timed exec>:1

File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/braket/experimental/algorithms/qc_qmc/qc_qmc.py:71, in qc_qmc(num_walkers, num_steps, dtau, quantum_evaluations_every_n_steps, trial, prop, trial_state_circuit, dev, max_pool)
     69 # parallelize with multiprocessing
     70 with mp.Pool(max_pool) as pool:
---> 71     results = list(pool.map(q_full_imag_time_evolution_wrapper, inputs))
     73 local_energies, weights, nums, denoms = map(np.array, zip(*results))
     75 energies = np.real(np.average(local_energies, weights=weights, axis=0))

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/pool.py:367, in Pool.map(self, func, iterable, chunksize)
    362 def map(self, func, iterable, chunksize=None):
    363     '''
    364     Apply `func` to each element in `iterable`, collecting the results
    365     in a list that is returned.
    366     '''
--> 367     return self._map_async(func, iterable, mapstar, chunksize).get()

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/pool.py:774, in ApplyResult.get(self, timeout)
    772     return self._value
    773 else:
--> 774     raise self._value

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/pool.py:540, in Pool._handle_tasks(taskqueue, put, outqueue, pool, cache)
    538     break
    539 try:
--> 540     put(task)
    541 except Exception as e:
    542     job, idx = task[:2]

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/connection.py:206, in _ConnectionBase.send(self, obj)
    204 self._check_closed()
    205 self._check_writable()
--> 206 self._send_bytes(_ForkingPickler.dumps(obj))

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/reduction.py:51, in ForkingPickler.dumps(cls, obj, protocol)
     48 @classmethod
     49 def dumps(cls, obj, protocol=None):
     50     buf = io.BytesIO()
---> 51     cls(buf, protocol).dump(obj)
     52     return buf.getbuffer()

TypeError: cannot pickle 'pennylane_lightning.lightning_qubit_ops.StateVectorC128' object

To reproduce
I created a new Braket NBI with default settings and ran the notebook.

Expected behavior
The code should execute without errors.

Screenshots or logs
If applicable, add screenshots or logs to help explain your problem.

System information
A description of your system. Please provide:

• Amazon Braket Python Algorithm Library version: 1.4.13
• Amazon Braket Python SDK version: 1.87.0
• Amazon Braket Python Schemas version: 1.22.0
• Amazon Braket Python Default Simulator version: 1.26.0
• PennyLane version: 0.35.1
• Python version: 3.10.14

Additional context

[licedric]

Error comes from sending the lightning.qubit device as an argument to pool.map for multiprocessing:

amazon-braket-algorithm-library/src/braket/experimental/algorithms/qc_qmc/qc_qmc.py

Line 71 in a1895e2

results = list(pool.map(q_full_imag_time_evolution_wrapper, inputs))

At some point the lightning.qubit device became unpicklable: https://discuss.pennylane.ai/t/qml-device-error-with-qiskit-aer/4556/7

Monkey-patching around this allows the notebook to run, so it's confirmed that this is the issue.