Executing differentiable circuits with adjoint and Sum expectation values raises an error

[mudit2812]

Describe the bug
I found this issue when attempting to run the Computing gradients in parallel with Amazon Braket PennyLane demo locally. When attempting to optimize a circuit measuring a LinearCombination, the execution raises a warning saying that adjoint differentation isn't supported for circuits with multi-term observables. This was a pennylane-side issue that is being addressed currently. However, after applying the PL-side fix, there is another error being raised. See details below. Note that I've only observed this failure for cases where the LinearCombination can be simplified. I've provided details in the "Additional context" section.

To reproduce
The following is a minimum working example where I observe the failure:

import pennylane as qml

n_wires = 5
device_arn = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"

dev = qml.device(
    "braket.aws.qubit",
    device_arn=device_arn,
    wires=n_wires,
    parallel=True,
    max_parallel=20,
    poll_timeout_seconds=30,
)

obs = qml.ops.LinearCombination([1.0, 2.0], [qml.X(0) @ qml.I(1), qml.Y(0) @ qml.X(1)])

@qml.qnode(dev)
def circuit(x):
    qml.RX(x, 0)
    return qml.expval(obs)

phi = qml.numpy.array(1.5, requires_grad=True)
print(circuit(phi))

Expected behavior
I expected the execution to give correct results.

Screenshots or logs
Running the above code block gives the following output:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[4], line 23
     20     return qml.expval(obs)
     22 phi = qml.numpy.array(1.5, requires_grad=True)
---> 23 print(circuit(phi))

File ~/repos/pennylane/pennylane/workflow/qnode.py:909, in QNode.__call__(self, *args, **kwargs)
    907 if qml.capture.enabled():
    908     return capture_qnode(self, *args, **kwargs)
--> 909 return self._impl_call(*args, **kwargs)

File ~/repos/pennylane/pennylane/workflow/qnode.py:885, in QNode._impl_call(self, *args, **kwargs)
    882 # Calculate the classical jacobians if necessary
    883 self._transform_program.set_classical_component(self, args, kwargs)
--> 885 res = qml.execute(
    886     (tape,),
    887     device=self.device,
    888     diff_method=self.diff_method,
    889     interface=interface,
    890     transform_program=self._transform_program,
    891     gradient_kwargs=self.gradient_kwargs,
    892     **self.execute_kwargs,
    893 )
    894 res = res[0]
    896 # convert result to the interface in case the qfunc has no parameters

File ~/repos/pennylane/pennylane/workflow/execution.py:232, in execute(tapes, device, diff_method, interface, transform_program, inner_transform, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, device_vjp, mcm_config, gradient_fn)
    229 if transform_program.is_informative:
    230     return post_processing(tapes)
--> 232 results = run(tapes, device, config, inner_transform)
    233 return post_processing(results)

File ~/repos/pennylane/pennylane/workflow/run.py:333, in run(tapes, device, config, inner_transform_program)
    330         params = tape.get_parameters(trainable_only=False)
    331         tape.trainable_params = qml.math.get_trainable_indices(params)
--> 333 results = ml_execute(tapes, execute_fn, jpc, device=device)
    334 return results

File ~/repos/pennylane/pennylane/workflow/interfaces/autograd.py:147, in autograd_execute(tapes, execute_fn, jpc, device)
    142 # pylint misidentifies autograd.builtins as a dict
    143 # pylint: disable=no-member
    144 parameters = autograd.builtins.tuple(
    145     [autograd.builtins.list(t.get_parameters()) for t in tapes]
    146 )
--> 147 return _execute(parameters, tuple(tapes), execute_fn, jpc)

File ~/.pyenv/versions/qml/lib/python3.10/site-packages/autograd/tracer.py:48, in primitive.<locals>.f_wrapped(*args, **kwargs)
     46     return new_box(ans, trace, node)
     47 else:
---> 48     return f_raw(*args, **kwargs)

File ~/repos/pennylane/pennylane/workflow/interfaces/autograd.py:183, in _execute(parameters, tapes, execute_fn, jpc)
    165 @autograd.extend.primitive
    166 def _execute(
    167     parameters,
   (...)
    170     jpc,
    171 ):  # pylint: disable=unused-argument
    172     """Autodifferentiable wrapper around a way of executing tapes.
    173 
    174     Args:
   (...)
    181 
    182     """
--> 183     return _to_autograd(execute_fn(tapes))

File ~/repos/pennylane/pennylane/workflow/jacobian_products.py:462, in DeviceDerivatives.execute_and_cache_jacobian(self, tapes)
    460 if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
    461     logger.debug("Forward pass called with %s", tapes)
--> 462 results, jac = self._dev_execute_and_compute_derivatives(tapes)
    463 self._results_cache[tapes] = results
    464 self._jacs_cache[tapes] = jac

File ~/repos/pennylane/pennylane/workflow/jacobian_products.py:427, in DeviceDerivatives._dev_execute_and_compute_derivatives(self, tapes)
    421 """
    422 Converts tapes to numpy before computing the the results and derivatives on the device.
    423 
    424 Dispatches between the two different device interfaces.
    425 """
    426 numpy_tapes, _ = qml.transforms.convert_to_numpy_parameters(tapes)
--> 427 return self._device.execute_and_compute_derivatives(numpy_tapes, self._execution_config)

File ~/repos/pennylane/pennylane/devices/modifiers/single_tape_support.py:62, in _make_execute_and_compute_derivatives.<locals>.execute_and_compute_derivatives(self, circuits, execution_config)
     60     is_single_circuit = True
     61     circuits = (circuits,)
---> 62 results, jacs = batch_execute_and_compute_derivatives(self, circuits, execution_config)
     63 return (results[0], jacs[0]) if is_single_circuit else (results, jacs)

File ~/repos/pennylane/pennylane/devices/legacy_facade.py:382, in LegacyDeviceFacade.execute_and_compute_derivatives(self, circuits, execution_config)
    380 first_shot = circuits[0].shots
    381 if all(t.shots == first_shot for t in circuits):
--> 382     return _set_shots(self._device, first_shot)(self._device.execute_and_gradients)(
    383         circuits, **execution_config.gradient_keyword_arguments
    384     )
    385 batched_res = tuple(
    386     self.execute_and_compute_derivatives((c,), execution_config) for c in circuits
    387 )
    388 return tuple(zip(*batched_res))

File ~/.pyenv/versions/3.10.15/lib/python3.10/contextlib.py:79, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
     76 @wraps(func)
     77 def inner(*args, **kwds):
     78     with self._recreate_cm():
---> 79         return func(*args, **kwds)

File ~/.pyenv/versions/qml/lib/python3.10/site-packages/braket/pennylane_plugin/braket_device.py:898, in BraketAwsQubitDevice.execute_and_gradients(self, circuits, **kwargs)
    896     new_res = self.execute(circuit, compute_gradient=False)
    897 else:
--> 898     results = self.execute(circuit, compute_gradient=True)
    899     new_res, new_jac = results
    900     # PennyLane expects the forward execution result to be a scalar
    901     # when it is accompanied by an adjoint gradient calculation

File ~/.pyenv/versions/qml/lib/python3.10/site-packages/braket/pennylane_plugin/braket_device.py:448, in BraketQubitDevice.execute(self, circuit, compute_gradient, **run_kwargs)
    442 self.check_validity(circuit.operations, circuit.observables)
    443 trainable = (
    444     BraketQubitDevice._get_trainable_parameters(circuit)
    445     if compute_gradient or self._parametrize_differentiable
    446     else {}
    447 )
--> 448 self._circuit = self._pl_to_braket_circuit(
    449     circuit,
    450     compute_gradient=compute_gradient,
    451     trainable_indices=frozenset(trainable.keys()),
    452     **run_kwargs,
    453 )
    454 if not isinstance(circuit.observables[0], MeasurementTransform):
    455     self._task = self._run_task(
    456         self._circuit, inputs={f"p_{k}": v for k, v in trainable.items()}
    457     )

File ~/.pyenv/versions/qml/lib/python3.10/site-packages/braket/pennylane_plugin/braket_device.py:259, in BraketQubitDevice._pl_to_braket_circuit(self, circuit, compute_gradient, trainable_indices, **run_kwargs)
    257     braket_circuit = Circuit().add_verbatim_box(braket_circuit)
    258 if compute_gradient:
--> 259     braket_circuit = self._apply_gradient_result_type(circuit, braket_circuit)
    260 elif not isinstance(circuit.measurements[0], MeasurementTransform):
    261     for measurement in circuit.measurements:

File ~/.pyenv/versions/qml/lib/python3.10/site-packages/braket/pennylane_plugin/braket_device.py:296, in BraketQubitDevice._apply_gradient_result_type(self, circuit, braket_circuit)
    292 else:
    293     targets = self.map_wires(pl_observable.wires).tolist()
    295 braket_circuit.add_result_type(
--> 296     get_adjoint_gradient_result_type(
    297         pl_observable,
    298         targets,
    299         self._braket_result_types,
    300         [f"p_{param_index}" for param_index in circuit.trainable_params],
    301     )
    302 )
    303 return braket_circuit

File ~/.pyenv/versions/qml/lib/python3.10/site-packages/braket/pennylane_plugin/translation.py:553, in get_adjoint_gradient_result_type(observable, targets, supported_result_types, parameters)
    549 braket_observable = _translate_observable(_flatten_observable(observable))
    550 braket_observable = (
    551     braket_observable.item() if hasattr(braket_observable, "item") else braket_observable
    552 )
--> 553 return AdjointGradient(observable=braket_observable, target=targets, parameters=parameters)

File ~/.pyenv/versions/qml/lib/python3.10/site-packages/braket/circuits/result_types.py:212, in AdjointGradient.__init__(self, observable, target, parameters)
    180 """Inits an `AdjointGradient`.
    181 
    182 Args:
   (...)
    209     >>> )
    210 """
    211 target_qubits = QubitSet(target if target is not None else observable.targets)
--> 212 super().__init__(
    213     ascii_symbols=[f"AdjointGradient({observable.ascii_symbols[0]})"] * len(target_qubits),
    214     observable=observable,
    215     target=target,
    216     parameters=parameters,
    217 )

File ~/.pyenv/versions/qml/lib/python3.10/site-packages/braket/circuits/result_type.py:293, in ObservableParameterResultType.__init__(self, ascii_symbols, observable, target, parameters)
    286 def __init__(
    287     self,
    288     ascii_symbols: list[str],
   (...)
    291     parameters: list[Union[str, FreeParameter]] | None = None,
    292 ):
--> 293     super().__init__(ascii_symbols, observable, target)
    295     self._parameters = (
    296         [(param.name if isinstance(param, FreeParameter) else param) for param in parameters]
    297         if parameters
    298         else parameters
    299     )
    301     """
    302     Args:
    303         ascii_symbols (list[str]): ASCII string symbols for the result type. This is used when
   (...)
    317             the number of `ascii_symbols` are not equal.
    318     """

File ~/.pyenv/versions/qml/lib/python3.10/site-packages/braket/circuits/result_type.py:222, in ObservableResultType.__init__(self, ascii_symbols, observable, target)
    220     for term_target, obs in zip(self._target, observable.summands):
    221         if obs.qubit_count != len(term_target):
--> 222             raise ValueError(
    223                 "Sum observable's target shape must be a nested list where each term's "
    224                 "target length is equal to the observable term's qubits count."
    225             )
    226 elif self._observable.qubit_count != len(self._target):
    227     raise ValueError(
    228         f"Observable's qubit count {self._observable.qubit_count} and "
    229         f"the size of the target qubit set {self._target} must be equal"
    230     )

ValueError: Sum observable's target shape must be a nested list where each term's target length is equal to the observable term's qubits count.

System information
A description of your system. Please provide:

• Amazon Braket Python PennyLane Plugin version: 1.31.0
• Amazon Braket Python SDK version: 1.88.3
• Amazon Braket Python Schemas version: 1.22.4
• Amazon Braket Python Default Simulator version: 1.26.2
• Python version: 3.10.15

Additional context
Based on my debugging, this is happening when the pennylane observable is being translated into a braket observable. In BraketQubitDevice._apply_gradient_result_type, we get the targets using the wires of the terms of the observable. However, in the subsequent call to get_adjoint_gradient_result_type (in translation.py), we simplify the pennylane observable (when calling _flatten_observable), but the targets are still the same as the unsimplified observable. This causes a mismatch when initializing AdjointGradient and raises the error seen above.

Note that if you tried running the example provided above with latest pennylane, the error will not occur (because of the warning mentioned in the bug description; a side effect of that warning is to change the diff method from adjoint to parameter-shift), so the mudit/0.40-fixes branch must be used. That branch will be merged into the v0.40 release candidate by end of week.