qml.sample() returns error on SV1 simulator

[cryotheta]

Describe the bug
qml.sample() throws error when run on SV1 simulator . The task is however executed on aws simulator and generates the json.

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-204-05637f991537> in <module>
      4     qml.Hadamard(wires=1)
      5     return qml.sample()
----> 6 circuit()

~/anaconda3/envs/Braket/lib/python3.7/site-packages/pennylane/qnode.py in __call__(self, *args, **kwargs)
    583             gradient_kwargs=self.gradient_kwargs,
    584             override_shots=override_shots,
--> 585             **self.execute_kwargs,
    586         )
    587 

~/anaconda3/envs/Braket/lib/python3.7/site-packages/pennylane/interfaces/batch/__init__.py in execute(tapes, device, gradient_fn, interface, mode, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform)
    411 
    412     res = _execute(
--> 413         tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=max_diff, mode=_mode
    414     )
    415 

~/anaconda3/envs/Braket/lib/python3.7/site-packages/pennylane/interfaces/batch/autograd.py in execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n, max_diff, mode)
     70         gradient_kwargs=gradient_kwargs,
     71         _n=_n,
---> 72         max_diff=max_diff,
     73     )[0]
     74 

~/anaconda3/envs/Braket/lib/python3.7/site-packages/autograd/tracer.py in f_wrapped(*args, **kwargs)
     46             return new_box(ans, trace, node)
     47         else:
---> 48             return f_raw(*args, **kwargs)
     49     f_wrapped.fun = f_raw
     50     f_wrapped._is_autograd_primitive = True

~/anaconda3/envs/Braket/lib/python3.7/site-packages/pennylane/interfaces/batch/autograd.py in _execute(parameters, tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n, max_diff)
    106     """
    107     with qml.tape.Unwrap(*tapes):
--> 108         res, jacs = execute_fn(tapes, **gradient_kwargs)
    109 
    110     for i, r in enumerate(res):

~/anaconda3/envs/Braket/lib/python3.7/site-packages/pennylane/interfaces/batch/__init__.py in wrapper(tapes, **kwargs)
    171         else:
    172             # execute all unique tapes that do not exist in the cache
--> 173             res = fn(execution_tapes.values(), **kwargs)
    174 
    175         final_res = []

~/anaconda3/envs/Braket/lib/python3.7/site-packages/pennylane/interfaces/batch/__init__.py in fn(tapes, **kwargs)
    123         def fn(tapes, **kwargs):  # pylint: disable=function-redefined
    124             tapes = [expand_fn(tape) for tape in tapes]
--> 125             return original_fn(tapes, **kwargs)
    126 
    127     @wraps(fn)

~/anaconda3/envs/Braket/lib/python3.7/contextlib.py in inner(*args, **kwds)
     72         def inner(*args, **kwds):
     73             with self._recreate_cm():
---> 74                 return func(*args, **kwds)
     75         return inner
     76 

~/anaconda3/envs/Braket/lib/python3.7/site-packages/braket/pennylane_plugin/braket_device.py in batch_execute(self, circuits, **run_kwargs)
    359     def batch_execute(self, circuits, **run_kwargs):
    360         if not self._parallel:
--> 361             return super().batch_execute(circuits)
    362 
    363         for circuit in circuits:

~/anaconda3/envs/Braket/lib/python3.7/site-packages/pennylane/_qubit_device.py in batch_execute(self, circuits)
    287             self.reset()
    288 
--> 289             res = self.execute(circuit)
    290             results.append(res)
    291 

~/anaconda3/envs/Braket/lib/python3.7/site-packages/braket/pennylane_plugin/braket_device.py in execute(self, circuit, **run_kwargs)
    218     def execute(self, circuit: CircuitGraph, **run_kwargs) -> np.ndarray:
    219         self.check_validity(circuit.operations, circuit.observables)
--> 220         self._circuit = self._pl_to_braket_circuit(circuit, **run_kwargs)
    221         self._task = self._run_task(self._circuit)
    222         braket_result = self._task.result()

~/anaconda3/envs/Braket/lib/python3.7/site-packages/braket/pennylane_plugin/braket_device.py in _pl_to_braket_circuit(self, circuit, **run_kwargs)
    151         for observable in circuit.observables:
    152             dev_wires = self.map_wires(observable.wires).tolist()
--> 153             translated = translate_result_type(observable, dev_wires, self._braket_result_types)
    154             if isinstance(translated, tuple):
    155                 for result_type in translated:

~/anaconda3/envs/Braket/lib/python3.7/site-packages/braket/pennylane_plugin/translation.py in translate_result_type(observable, targets, supported_result_types)
    353         raise NotImplementedError(f"Return type {return_type} unsupported for Hamiltonian")
    354 
--> 355     braket_observable = _translate_observable(observable)
    356     if return_type is ObservableReturnTypes.Expectation:
    357         return Expectation(braket_observable, targets)

~/anaconda3/envs/Braket/lib/python3.7/functools.py in wrapper(*args, **kw)
    838                             '1 positional argument')
    839 
--> 840         return dispatch(args[0].__class__)(*args, **kw)
    841 
    842     funcname = getattr(func, '__name__', 'singledispatch function')

~/anaconda3/envs/Braket/lib/python3.7/site-packages/braket/pennylane_plugin/translation.py in _translate_observable(observable)
    366 @singledispatch
    367 def _translate_observable(observable):
--> 368     raise TypeError(f"Unsupported observable: {observable}")
    369 
    370 

TypeError: Unsupported observable: sample(wires=[])

To reproduce

arn = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
dev = qml.device("braket.aws.qubit",device_arn=arn, wires=2,shots=1024)
@qml.qnode(dev)
def circuit():
    qml.PauliX(wires=0)
    qml.Hadamard(wires=1)
    return qml.sample()
circuit()

Expected behavior
Should return tensors containing samples.
tensor([[1, 0, 0, 0, 0, 0],
[1, 1, 0, 0, 0, 0],
[1, 0, 0, 0, 0, 0],
...,
[1, 0, 0, 0, 0, 0],
[1, 1, 0, 0, 0, 0],
[1, 0, 0, 0, 0, 0]], requires_grad=True)
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 PennyLane Plugin version:
• 1.6.5
• Amazon Braket Python SDK version:
• 1.24.0
• Amazon Braket Python Schemas version:
• 1.8.0.dev0
• Amazon Braket Python Default Simulator version:
• 1.3.0
• Python version: 3.7

Additional context

[krneta]

Hi @cryotheta ,

Thanks for bringing this to our attention. In your code, please specify the observable you wish to sample, for example:

arn = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
dev = qml.device("braket.aws.qubit",device_arn=arn, wires=2,shots=1024)
@qml.qnode(dev)
def circuit():
    qml.PauliX(wires=0)
    qml.Hadamard(wires=1)
    return qml.sample(qml.PauliZ(1))
circuit()

You can pick from: PauliX, PauliY, PauliZ, or Hadamard.

Please let me know if this answers your question. We'll look into if we can make the error more clear.

[cryotheta]

I want to measure the entire circuit(all wires) in the computational basis. thus, as per the pennylane documentation qml.sample() with None as input should have done the trick. (https://pennylane.readthedocs.io/en/stable/code/api/pennylane.sample.html)
Specifying qml.sample(PauliZ(1)) for all wires will drastically increase the runtime and cost. Can you recommend a solution?

[krneta]

Hi @cryotheta ,

If you just call sample multiple times, would that still dramatically increase the runtime and cost?

For example:

return [qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1))]

[virajvchaudhari]

Hi @cryotheta, was your issue resolved or are you still looking for help on this?

[virajvchaudhari]

Hi @cryotheta,

Based on the earlier suggestions made in response to your query, I will close this issue.
If your questions are still unanswered, please feel free to re-open this issue.