Braket plugin fails on 16 or more qubits

[licedric]

Prerequisites

Before raising this issue, I have already checked that I am:

• running the latest version
• made sure that this issue has not already been filed

Describe the bug

The Braket plugin fails when running circuits of 16 or more qubits.

To Reproduce

Steps to reproduce the behavior:
Run this tutorial notebook in the documentation. This gives an error with the error message ValidationException: An error occurred (ValidationException) when calling the CreateQuantumTask operation: Probability result type must be used with 15 or fewer qubits.. See additional context for full stack trace.

Expected behavior

The notebook should run without errors on 16 or more qubits.

Additional context

Two things are in play here:

1. Braket limits the Probability result type to use no more than 15 qubits, since the result type is returned a dense vector and is mainly meant to be used for marginal probability distributions. Unfortunately this 15-qubit target limit is not yet in the documentation.
2. OpenQAOA uses the Probability result type in this line. This is used for all circuits and targeting all qubits, so it will fail when the circuit has 16 or more qubits.

It actually doesn't look like OpenQAOA uses the Probability result at all, and just looks at the measurement counts (see this line). Perhaps we can just remove the use of the Probability result type?

Full stack trace:

---------------------------------------------------------------------------
ValidationException                       Traceback (most recent call last)
Cell In[21], line 1
----> 1 q.optimize()

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/openqaoa/algorithms/qaoa/qaoa_workflow.py:383, in QAOA.optimize(self, verbose)
    380 # timestamp for the start of the optimization
    381 self.header["execution_time_start"] = generate_timestamp()
--> 383 self.optimizer.optimize()
    384 # TODO: result and qaoa_result will differ
    385 self.result = self.optimizer.qaoa_result

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/openqaoa/optimizers/training_vqa.py:516, in ScipyOptimizer.optimize(self)
    510     print(
    511         "The optimization has been terminated early. Most likely due to a connection error."
    512         "You can retrieve results from the optimization runs that were completed"
    513         "through the .result method."
    514     )
    515 except Exception as e:
--> 516     raise e
    517 finally:
    518     self.results_dictionary()

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/openqaoa/optimizers/training_vqa.py:499, in ScipyOptimizer.optimize(self)
    487             result = minimize(
    488                 self.optimize_this,
    489                 x0=self.initial_params,
   (...)
    496                 bounds=self.bounds,
    497             )
    498     else:
--> 499         result = minimize(
    500             self.optimize_this,
    501             x0=self.initial_params,
    502             method=self.method,
    503             tol=self.tol,
    504             constraints=self.constraints,
    505             options=self.options,
    506             bounds=self.bounds,
    507         )
    508 except ConnectionError as e:
    509     print(e, "\n")

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/scipy/optimize/_minimize.py:737, in minimize(fun, x0, args, method, jac, hess, hessp, bounds, constraints, tol, callback, options)
    734     res = _minimize_tnc(fun, x0, args, jac, bounds, callback=callback,
    735                         **options)
    736 elif meth == 'cobyla':
--> 737     res = _minimize_cobyla(fun, x0, args, constraints, callback=callback,
    738                            bounds=bounds, **options)
    739 elif meth == 'cobyqa':
    740     res = _minimize_cobyqa(fun, x0, args, bounds, constraints, callback,
    741                            **options)

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/scipy/optimize/_cobyla_py.py:35, in synchronized.<locals>.wrapper(*args, **kwargs)
     32 @functools.wraps(func)
     33 def wrapper(*args, **kwargs):
     34     with _module_lock:
---> 35         return func(*args, **kwargs)

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/scipy/optimize/_cobyla_py.py:278, in _minimize_cobyla(fun, x0, args, constraints, rhobeg, tol, maxiter, disp, catol, callback, bounds, **unknown_options)
    275 def _jac(x, *args):
    276     return None
--> 278 sf = _prepare_scalar_function(fun, x0, args=args, jac=_jac)
    280 def calcfc(x, con):
    281     f = sf.fun(x)

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/scipy/optimize/_optimize.py:288, in _prepare_scalar_function(fun, x0, jac, args, bounds, epsilon, finite_diff_rel_step, hess)
    284     bounds = (-np.inf, np.inf)
    286 # ScalarFunction caches. Reuse of fun(x) during grad
    287 # calculation reduces overall function evaluations.
--> 288 sf = ScalarFunction(fun, x0, args, grad, hess,
    289                     finite_diff_rel_step, bounds, epsilon=epsilon)
    291 return sf

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/scipy/optimize/_differentiable_functions.py:222, in ScalarFunction.__init__(self, fun, x0, args, grad, hess, finite_diff_rel_step, finite_diff_bounds, epsilon)
    219     finite_diff_options["as_linear_operator"] = True
    221 # Initial function evaluation
--> 222 self._update_fun()
    224 # Initial gradient evaluation
    225 self._wrapped_grad, self._ngev = _wrapper_grad(
    226     grad,
    227     fun=self._wrapped_fun,
    228     args=args,
    229     finite_diff_options=finite_diff_options
    230 )

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/scipy/optimize/_differentiable_functions.py:294, in ScalarFunction._update_fun(self)
    292 def _update_fun(self):
    293     if not self.f_updated:
--> 294         fx = self._wrapped_fun(self.x)
    295         if fx < self._lowest_f:
    296             self._lowest_x = self.x

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/scipy/optimize/_differentiable_functions.py:20, in _wrapper_fun.<locals>.wrapped(x)
     16 ncalls[0] += 1
     17 # Send a copy because the user may overwrite it.
     18 # Overwriting results in undefined behaviour because
     19 # fun(self.x) will change self.x, with the two no longer linked.
---> 20 fx = fun(np.copy(x), *args)
     21 # Make sure the function returns a true scalar
     22 if not np.isscalar(fx):

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/openqaoa/optimizers/training_vqa.py:232, in OptimizeVQA.optimize_this(self, x, n_shots)
    229     save_parameter("param_log", deepcopy(x))
    231 n_shots_dict = {"n_shots": n_shots} if n_shots else {}
--> 232 callback_cost = self.vqa.expectation(self.variational_params, **n_shots_dict)
    234 log_dict.update({"cost": callback_cost})
    236 current_eval = self.log.func_evals.best[0]

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/openqaoa/utilities.py:252, in round_value.<locals>.wrapper(*args, **kwargs)
    251 def wrapper(*args, **kwargs):
--> 252     values = function(*args, **kwargs)
    253     if isinstance(values, dict):
    254         return {k: round(v, PRECISION) for k, v in values.items()}

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/openqaoa/backends/basebackend.py:282, in QAOABaseBackend.expectation(self, params, n_shots)
    264 @round_value
    265 def expectation(self, params: QAOAVariationalBaseParams, n_shots=None) -> float:
    266     """
    267     Compute the expectation value w.r.t the Cost Hamiltonian
    268 
   (...)
    280         Expectation value of cost operator wrt to quantum state produced by QAOA circuit
    281     """
--> 282     counts = self.get_counts(params, n_shots)
    283     cost = cost_function(
    284         counts, self.qaoa_descriptor.cost_hamiltonian, self.cvar_alpha
    285     )
    286     return cost

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/openqaoa_braket/backends/qaoa_braket_qpu.py:220, in QAOAAWSQPUBackend.get_counts(self, params, n_shots)
    217 max_job_retries = 5
    219 while job_state == False:
--> 220     job = self.device.backend_device.run(
    221         circuit,
    222         (self.device.s3_bucket_name, self.device.folder_name),
    223         shots=n_shots,
    224         disable_qubit_rewiring=self.disable_qubit_rewiring,
    225     )
    227     try:
    228         self.job_id = job.id

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/braket/aws/aws_device.py:208, in AwsDevice.run(self, task_specification, s3_destination_folder, shots, poll_timeout_seconds, poll_interval_seconds, inputs, gate_definitions, reservation_arn, *aws_quantum_task_args, **aws_quantum_task_kwargs)
    206 if self._noise_model:
    207     task_specification = self._apply_noise_model_to_circuit(task_specification)
--> 208 return AwsQuantumTask.create(
    209     self._aws_session,
    210     self._arn,
    211     task_specification,
    212     s3_destination_folder
    213     or (
    214         AwsSession.parse_s3_uri(os.environ.get("AMZN_BRAKET_TASK_RESULTS_S3_URI"))
    215         if "AMZN_BRAKET_TASK_RESULTS_S3_URI" in os.environ
    216         else None
    217     )
    218     or (self._aws_session.default_bucket(), "tasks"),
    219     shots if shots is not None else self._default_shots,
    220     poll_timeout_seconds=poll_timeout_seconds,
    221     poll_interval_seconds=poll_interval_seconds or self._poll_interval_seconds,
    222     inputs=inputs,
    223     gate_definitions=gate_definitions,
    224     reservation_arn=reservation_arn,
    225     *aws_quantum_task_args,
    226     **aws_quantum_task_kwargs,
    227 )

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/braket/aws/aws_quantum_task.py:214, in AwsQuantumTask.create(aws_session, device_arn, task_specification, s3_destination_folder, shots, device_parameters, disable_qubit_rewiring, tags, inputs, gate_definitions, quiet, reservation_arn, *args, **kwargs)
    209     if unbounded_parameters := param_names - set(inputs.keys()):
    210         raise ValueError(
    211             f"Cannot execute circuit with unbound parameters: {unbounded_parameters}"
    212         )
--> 214 return _create_internal(
    215     task_specification,
    216     aws_session,
    217     create_task_kwargs,
    218     device_arn,
    219     device_parameters or {},
    220     disable_qubit_rewiring,
    221     inputs,
    222     gate_definitions=gate_definitions,
    223     quiet=quiet,
    224     *args,
    225     **kwargs,
    226 )

File ~/.pyenv/versions/3.10.14/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 ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/braket/aws/aws_quantum_task.py:730, in _(circuit, aws_session, create_task_kwargs, device_arn, device_parameters, disable_qubit_rewiring, inputs, gate_definitions, *args, **kwargs)
    721     openqasm_program = OpenQASMProgram(
    722         source=openqasm_program.source,
    723         inputs=inputs_copy,
    724     )
    726 create_task_kwargs |= {
    727     "action": openqasm_program.json(),
    728     "deviceParameters": final_device_parameters.json(exclude_none=True),
    729 }
--> 730 task_arn = aws_session.create_quantum_task(**create_task_kwargs)
    731 return AwsQuantumTask(task_arn, aws_session, *args, **kwargs)

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/braket/aws/aws_session.py:269, in AwsSession.create_quantum_task(self, **boto3_kwargs)
    267 if job_token:
    268     boto3_kwargs["jobToken"] = job_token
--> 269 response = self.braket_client.create_quantum_task(**boto3_kwargs)
    270 broadcast_event(
    271     _TaskCreationEvent(
    272         arn=response["quantumTaskArn"],
   (...)
    276     )
    277 )
    278 return response["quantumTaskArn"]

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/botocore/client.py:569, in ClientCreator._create_api_method.<locals>._api_call(self, *args, **kwargs)
    565     raise TypeError(
    566         f"{py_operation_name}() only accepts keyword arguments."
    567     )
    568 # The "self" in this scope is referring to the BaseClient.
--> 569 return self._make_api_call(operation_name, kwargs)

File ~/.pyenv/versions/3.10.14/lib/python3.10/site-packages/botocore/client.py:1023, in BaseClient._make_api_call(self, operation_name, api_params)
   1019     error_code = error_info.get("QueryErrorCode") or error_info.get(
   1020         "Code"
   1021     )
   1022     error_class = self.exceptions.from_code(error_code)
-> 1023     raise error_class(parsed_response, operation_name)
   1024 else:
   1025     return parsed_response

ValidationException: An error occurred (ValidationException) when calling the CreateQuantumTask operation: Probability result type must be used with 15 or fewer qubits.