From b12f02c6ba14f3d74113407c236bbd527de5131f Mon Sep 17 00:00:00 2001 From: albi3ro Date: Fri, 16 Aug 2024 10:41:54 -0400 Subject: [PATCH 1/5] Improvements to jitting --- doc/releases/changelog-dev.md | 11 ++ pennylane/measurements/classical_shadow.py | 14 +-- pennylane/measurements/expval.py | 7 +- pennylane/measurements/measurements.py | 64 +++------- pennylane/measurements/mutual_info.py | 7 +- pennylane/measurements/probs.py | 13 +- pennylane/measurements/purity.py | 7 +- pennylane/measurements/sample.py | 24 +--- pennylane/measurements/state.py | 21 +--- pennylane/measurements/var.py | 7 +- pennylane/measurements/vn_entropy.py | 7 +- pennylane/tape/qscript.py | 39 +++--- pennylane/workflow/interfaces/jax_jit.py | 60 +++++---- pennylane/workflow/qnode.py | 7 +- tests/interfaces/test_jax_jit.py | 14 +++ tests/measurements/test_measurements.py | 47 ++++++++ tests/tape/test_qscript.py | 134 +-------------------- tests/tape/test_tape.py | 48 -------- tests/test_qnode.py | 4 +- tests/test_qnode_legacy.py | 4 +- 20 files changed, 179 insertions(+), 360 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 4dd6effac75..bb023b13e79 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -27,6 +27,10 @@

Improvements 🛠

+* Counts measurements with `all_outcomes=True` can now be used with jax jitting. Measurements + broadcasted across all available wires (`qml.probs()`) can now be used with jit and devices that + allow variable numbers of wires (`qml.device('default.qubit')`). + * Mid-circuit measurements can now be captured with `qml.capture` enabled. [(#6015)](https://github.com/PennyLaneAI/pennylane/pull/6015) @@ -226,6 +230,13 @@

Breaking changes 💔

+* `MeasurementProcess.shape(shots: Shots, device:Device)` is now + `MeasurementProcess.shape(shots: Optional[int], num_device_wires:int = 0)`. This has been done to allow + jitting when a measurement is broadcasted across all available wires, but the device does not specify wires. + +* If the shape of a probability measurement is affected by a `Device.cutoff` property, it will no longer work with + jitting. + * `GlobalPhase` is considered non-differentiable with tape transforms. As a consequence, `qml.gradients.finite_diff` and `qml.gradients.spsa_grad` no longer support differentiation of `GlobalPhase` with state-based outputs. diff --git a/pennylane/measurements/classical_shadow.py b/pennylane/measurements/classical_shadow.py index 7ac08e7a758..6c502858cf8 100644 --- a/pennylane/measurements/classical_shadow.py +++ b/pennylane/measurements/classical_shadow.py @@ -450,9 +450,9 @@ def _abstract_eval( ) -> tuple: return (2, shots, n_wires), np.int8 - def shape(self, device, shots): # pylint: disable=unused-argument + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int, int, int]: # otherwise, the return type requires a device - if not shots: + if shots is None: raise MeasurementShapeError( "Shots must be specified to obtain the shape of a classical " "shadow measurement process." @@ -460,7 +460,7 @@ def shape(self, device, shots): # pylint: disable=unused-argument # the first entry of the tensor represents the measured bits, # and the second indicate the indices of the unitaries used - return (2, shots.total_shots, len(self.wires)) + return (2, shots, len(self.wires)) def __copy__(self): return self.__class__( @@ -559,12 +559,8 @@ def numeric_type(self): def return_type(self): return ShadowExpval - def shape(self, device, shots): - is_single_op = isinstance(self.H, Operator) - if not shots.has_partitioned_shots: - return () if is_single_op else (len(self.H),) - base = () if is_single_op else (len(self.H),) - return (base,) * sum(s.copies for s in shots.shot_vector) + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple: + return () if isinstance(self.H, Operator) else (len(self.H),) @property def wires(self): diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py index 95499fe9ea8..1a95f9e7ce0 100644 --- a/pennylane/measurements/expval.py +++ b/pennylane/measurements/expval.py @@ -98,11 +98,8 @@ class ExpectationMP(SampleMeasurement, StateMeasurement): def numeric_type(self): return float - def shape(self, device, shots): - if not shots.has_partitioned_shots: - return () - num_shot_elements = sum(s.copies for s in shots.shot_vector) - return tuple(() for _ in range(num_shot_elements)) + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple: + return () def process_samples( self, diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py index b6d41dd097e..ab28f487415 100644 --- a/pennylane/measurements/measurements.py +++ b/pennylane/measurements/measurements.py @@ -17,7 +17,6 @@ and measurement samples using AnnotatedQueues. """ import copy -import functools from abc import ABC, abstractmethod from collections.abc import Sequence from enum import Enum @@ -30,8 +29,6 @@ from pennylane.typing import TensorLike from pennylane.wires import Wires -from .shots import Shots - # ============================================================================= # ObservableReturnTypes types # ============================================================================= @@ -289,58 +286,35 @@ def numeric_type(self) -> type: f"The numeric type of the measurement {self.__class__.__name__} is not defined." ) - def shape(self, device, shots: Shots) -> tuple: - """The expected output shape of the MeasurementProcess. - - Note that the output shape is dependent on the shots or device when: - - * The measurement type is either ``_Probability``, ``_State`` (from :func:`.state`) or - ``_Sample``; - * The shot vector was defined. - - For example, assuming a device with ``shots=None``, expectation values - and variances define ``shape=(,)``, whereas probabilities in the qubit - model define ``shape=(2**num_wires)`` where ``num_wires`` is the - number of wires the measurement acts on. + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int, ...]: + """Calculate the shape of the result object tensor. Args: - device (pennylane.Device): a PennyLane device to use for determining the shape - shots (~.Shots): object defining the number and batches of shots + shots (Optional[int]) = None: the number of shots used execute the circuit. ``None`` + indicates an analytic simulation. Shot vectors are handled by calling this method + multiple times. + num_device_wires (int)=0 : The number of wires that will be used if the measurement is + broadcasted across all available wires (``len(mp.wires) == 0``). If the device + itself doesn't provide a number of wires, the number of tape wires will be provided + here instead: Returns: - tuple: the output shape + tuple[int,...]: An arbitrary length tuple of ints. May be an empty tuple. + + >>> qml.probs(wires=(0,1)).shape() + (4,) + >>> qml.sample(wires=(0,1)).shape(shots=50) + (50, 2) + >>> qml.state().shape(num_device_wires=4) + (16,) + >>> qml.expval(qml.Z(0)).shape() + () - Raises: - QuantumFunctionError: the return type of the measurement process is - unrecognized and cannot deduce the numeric type """ raise qml.QuantumFunctionError( f"The shape of the measurement {self.__class__.__name__} is not defined" ) - @staticmethod - @functools.lru_cache() - def _get_num_basis_states(num_wires, device): - """Auxiliary function to determine the number of basis states given the - number of systems and a quantum device. - - This function is meant to be used with the Probability measurement to - determine how many outcomes there will be. With qubit based devices - we'll have two outcomes for each subsystem. With continuous variable - devices that impose a Fock cutoff the number of basis states per - subsystem equals the cutoff value. - - Args: - num_wires (int): the number of qubits/qumodes - device (pennylane.Device): a PennyLane device - - Returns: - int: the number of basis states - """ - cutoff = getattr(device, "cutoff", None) - base = 2 if cutoff is None else cutoff - return base**num_wires - @qml.QueuingManager.stop_recording() def diagonalizing_gates(self): """Returns the gates that diagonalize the measured wires such that they diff --git a/pennylane/measurements/mutual_info.py b/pennylane/measurements/mutual_info.py index 95c71f19615..18335fe4306 100644 --- a/pennylane/measurements/mutual_info.py +++ b/pennylane/measurements/mutual_info.py @@ -154,11 +154,8 @@ def map_wires(self, wire_map: dict): ] return new_measurement - def shape(self, device, shots): - if not shots.has_partitioned_shots: - return () - num_shot_elements = sum(s.copies for s in shots.shot_vector) - return tuple(() for _ in range(num_shot_elements)) + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple: + return () def process_state(self, state: Sequence[complex], wire_order: Wires): state = qml.math.dm_from_state_vector(state) diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py index 74e62bf122d..5bc71b5ed88 100644 --- a/pennylane/measurements/probs.py +++ b/pennylane/measurements/probs.py @@ -163,16 +163,9 @@ def _abstract_eval(cls, n_wires=None, has_eigvals=False, shots=None, num_device_ def numeric_type(self): return float - def shape(self, device, shots): - num_shot_elements = ( - sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1 - ) - len_wires = len(self.wires) - if len_wires == 0: - len_wires = len(device.wires) if device.wires else 0 - dim = self._get_num_basis_states(len_wires, device) - - return (dim,) if num_shot_elements == 1 else tuple((dim,) for _ in range(num_shot_elements)) + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int]: + len_wires = len(self.wires) if self.wires else num_device_wires + return (2**len_wires,) def process_samples( self, diff --git a/pennylane/measurements/purity.py b/pennylane/measurements/purity.py index e071f05dd1b..83e4166cbc9 100644 --- a/pennylane/measurements/purity.py +++ b/pennylane/measurements/purity.py @@ -85,11 +85,8 @@ def return_type(self): def numeric_type(self): return float - def shape(self, device, shots): - if not shots.has_partitioned_shots: - return () - num_shot_elements = sum(s.copies for s in shots.shot_vector) - return tuple(() for _ in range(num_shot_elements)) + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple: + return () def process_state(self, state: Sequence[complex], wire_order: Wires): wire_map = dict(zip(wire_order, list(range(len(wire_order))))) diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py index b50c7f0a08d..e341cbe75f0 100644 --- a/pennylane/measurements/sample.py +++ b/pennylane/measurements/sample.py @@ -220,7 +220,7 @@ def numeric_type(self): return int return float - def shape(self, device, shots): + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple: if not shots: raise MeasurementShapeError( "Shots are required to obtain the shape of the measurement " @@ -230,25 +230,13 @@ def shape(self, device, shots): num_values_per_shot = 1 # one single eigenvalue else: # one value per wire - num_values_per_shot = len(self.wires) if len(self.wires) > 0 else len(device.wires) - - def _single_int_shape(shot_val, num_values): - # singleton dimensions, whether in shot val or num_wires are squeezed away - inner_shape = [] - if shot_val != 1: - inner_shape.append(shot_val) - if num_values != 1: - inner_shape.append(num_values) - return tuple(inner_shape) - - if not shots.has_partitioned_shots: - return _single_int_shape(shots.total_shots, num_values_per_shot) + num_values_per_shot = len(self.wires) if len(self.wires) > 0 else num_device_wires shape = [] - for s in shots.shot_vector: - for _ in range(s.copies): - shape.append(_single_int_shape(s.shots, num_values_per_shot)) - + if shots != 1: + shape.append(shots) + if num_values_per_shot != 1: + shape.append(num_values_per_shot) return tuple(shape) def process_samples( diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py index b6993427179..5b8baecf8b7 100644 --- a/pennylane/measurements/state.py +++ b/pennylane/measurements/state.py @@ -159,12 +159,9 @@ def _abstract_eval( def numeric_type(self): return complex - def shape(self, device, shots): - num_shot_elements = ( - sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1 - ) - dim = 2 ** len(self.wires) if self.wires else 2 ** len(device.wires) - return (dim,) if num_shot_elements == 1 else tuple((dim,) for _ in range(num_shot_elements)) + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int]: + num_wires = len(self.wires) if self.wires else num_device_wires + return (2**num_wires,) def process_state(self, state: Sequence[complex], wire_order: Wires): # pylint:disable=redefined-outer-name @@ -232,17 +229,9 @@ def _abstract_eval( shape = (2**n_wires, 2**n_wires) return shape, complex - def shape(self, device, shots): - num_shot_elements = ( - sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1 - ) - + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int, int]: dim = 2 ** len(self.wires) - return ( - (dim, dim) - if num_shot_elements == 1 - else tuple((dim, dim) for _ in range(num_shot_elements)) - ) + return (dim, dim) def process_state(self, state: Sequence[complex], wire_order: Wires): # pylint:disable=redefined-outer-name diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py index 21c9bf1345f..ab0448ffc6d 100644 --- a/pennylane/measurements/var.py +++ b/pennylane/measurements/var.py @@ -90,11 +90,8 @@ class VarianceMP(SampleMeasurement, StateMeasurement): def numeric_type(self): return float - def shape(self, device, shots): - if not shots.has_partitioned_shots: - return () - num_shot_elements = sum(s.copies for s in shots.shot_vector) - return tuple(() for _ in range(num_shot_elements)) + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple: + return () def process_samples( self, diff --git a/pennylane/measurements/vn_entropy.py b/pennylane/measurements/vn_entropy.py index 1589cee7b90..79c1a875eae 100644 --- a/pennylane/measurements/vn_entropy.py +++ b/pennylane/measurements/vn_entropy.py @@ -113,11 +113,8 @@ def return_type(self): def numeric_type(self): return float - def shape(self, device, shots): - if not shots.has_partitioned_shots: - return () - num_shot_elements = sum(s.copies for s in shots.shot_vector) - return tuple(() for _ in range(num_shot_elements)) + def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple: + return () def process_state(self, state: Sequence[complex], wire_order: Wires): state = qml.math.dm_from_state_vector(state) diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py index 271929ad8c8..d21a3a09b51 100644 --- a/pennylane/tape/qscript.py +++ b/pennylane/tape/qscript.py @@ -763,7 +763,7 @@ def shape(self, device): dependent on the device used for execution. Args: - device (pennylane.Device): the device that will be used for the script execution + device (pennylane.devices.Device): the device that will be used for the script execution Returns: Union[tuple[int], tuple[tuple[int]]]: the output shape(s) of the quantum script result @@ -781,27 +781,22 @@ def shape(self, device): >>> qs.shape(dev) ((4,), (), (4,)) """ - shots = self.shots - # even with the legacy device interface, the shots on the tape will agree with the shots used by the device for the execution - - if len(shots.shot_vector) > 1 and self.batch_size is not None: - raise NotImplementedError( - "Parameter broadcasting when using a shot vector is not supported yet." - ) - - shapes = tuple(meas_process.shape(device, shots) for meas_process in self.measurements) - - if self.batch_size is not None: - shapes = tuple((self.batch_size,) + shape for shape in shapes) - - if len(shapes) == 1: - return shapes[0] - - if shots.num_copies > 1: - # put the shot vector axis before the measurement axis - shapes = tuple(zip(*shapes)) - - return shapes + num_device_wires = len(device.wires) if device.wires else len(self.wires) + + def get_shape(mp, _shots): + # depends on num_device_wires and self.batch_size from closure + standard_shape = mp.shape(shots=_shots, num_device_wires=num_device_wires) + if self.batch_size: + return (self.batch_size, *standard_shape) + return standard_shape + + shape = [] + for s in self.shots if self.shots else [None]: + shots_shape = tuple(get_shape(mp, s) for mp in self.measurements) + shots_shape = shots_shape[0] if len(shots_shape) == 1 else tuple(shots_shape) + shape.append(shots_shape) + + return tuple(shape) if self.shots.has_partitioned_shots else shape[0] @property def numeric_type(self): diff --git a/pennylane/workflow/interfaces/jax_jit.py b/pennylane/workflow/interfaces/jax_jit.py index 552748748f2..f00028f07a6 100644 --- a/pennylane/workflow/interfaces/jax_jit.py +++ b/pennylane/workflow/interfaces/jax_jit.py @@ -58,9 +58,8 @@ def _to_jax(result: qml.typing.ResultBatch) -> qml.typing.ResultBatch: ResultBatch: a nested structure of tuples, and jax arrays """ - # jax-jit not compatible with counts - # if isinstance(result, dict): - # return result + if isinstance(result, dict): + return {key: jnp.array(value) for key, value in result.items()} if isinstance(result, (list, tuple)): return tuple(_to_jax(r) for r in result) return jnp.array(result) @@ -86,23 +85,40 @@ def _jax_dtype(m_type): return jnp.dtype(m_type) +def _get_counts_shape(mp: "qml.measurements.CountsMP", num_device_wires=0): + num_wires = len(mp.wires) if mp.wires else num_device_wires + outcome_counts = {} + binary_pattern = "{0:0" + str(num_wires) + "b}" + for outcome in range(2**num_wires): + outcome_binary = binary_pattern.format(outcome) + outcome_counts[outcome_binary] = jax.core.ShapedArray((), _jax_dtype(int)) + + return outcome_counts + + def _result_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"): """Auxiliary function for creating the shape and dtype object structure given a tape.""" - shape = tape.shape(device) - if len(tape.measurements) == 1: - m_dtype = _jax_dtype(tape.measurements[0].numeric_type) - if tape.shots.has_partitioned_shots: - return tuple(jax.ShapeDtypeStruct(s, m_dtype) for s in shape) - return jax.ShapeDtypeStruct(tuple(shape), m_dtype) - - tape_dtype = tuple(_jax_dtype(m.numeric_type) for m in tape.measurements) - if tape.shots.has_partitioned_shots: - return tuple( - tuple(jax.ShapeDtypeStruct(tuple(s), d) for s, d in zip(si, tape_dtype)) for si in shape - ) - return tuple(jax.ShapeDtypeStruct(tuple(s), d) for s, d in zip(shape, tape_dtype)) + num_device_wires = len(device.wires) if device.wires else len(tape.wires) + + def struct(mp, shots): + # depends on num_device_wires and tape.batch_size from closure + if isinstance(mp, qml.measurements.CountsMP): + return _get_counts_shape(mp) + mp_shape = mp.shape(shots=shots, num_device_wires=num_device_wires) + if tape.batch_size: + mp_shape = (tape.batch_size, *mp_shape) + return jax.ShapeDtypeStruct(mp_shape, _jax_dtype(mp.numeric_type)) + + shape = [] + for s in tape.shots if tape.shots else [None]: + shots_shape = tuple(struct(mp, s) for mp in tape.measurements) + + shots_shape = shots_shape[0] if len(shots_shape) == 1 else tuple(shots_shape) + shape.append(shots_shape) + + return tuple(shape) if tape.shots.has_partitioned_shots else shape[0] def _jac_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"): @@ -234,11 +250,7 @@ def jax_jit_jvp_execute(tapes, execute_fn, jpc, device): """ - if any( - m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts) - for t in tapes - for m in t.measurements - ): + if any(m.return_type == qml.measurements.Counts for t in tapes for m in t.measurements): # Obtaining information about the shape of the Counts measurements is # not implemented and is required for the callback logic raise NotImplementedError("The JAX-JIT interface doesn't support qml.counts.") @@ -264,11 +276,7 @@ def jax_jit_vjp_execute(tapes, execute_fn, jpc, device=None): the returned tuple corresponds in order to the provided tapes. """ - if any( - m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts) - for t in tapes - for m in t.measurements - ): + if any(m.return_type == qml.measurements.Counts for t in tapes for m in t.measurements): # Obtaining information about the shape of the Counts measurements is # not implemented and is required for the callback logic raise NotImplementedError("The JAX-JIT interface doesn't support qml.counts.") diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py index 92b1818cff4..54d987c126f 100644 --- a/pennylane/workflow/qnode.py +++ b/pennylane/workflow/qnode.py @@ -27,7 +27,7 @@ from pennylane import Device from pennylane.debugging import pldb_device_manager from pennylane.logging import debug_logger -from pennylane.measurements import CountsMP, MidMeasureMP, Shots +from pennylane.measurements import MidMeasureMP, Shots from pennylane.tape import QuantumScript, QuantumTape from pennylane.transforms.core import TransformContainer, TransformDispatcher, TransformProgram @@ -1024,11 +1024,6 @@ def construct(self, args, kwargs): # pylint: disable=too-many-branches params = self.tape.get_parameters(trainable_only=False) self.tape.trainable_params = qml.math.get_trainable_indices(params) - if any(isinstance(m, CountsMP) for m in self.tape.measurements) and any( - qml.math.is_abstract(a) for a in args - ): - raise qml.QuantumFunctionError("Can't JIT a quantum function that returns counts.") - if isinstance(self._qfunc_output, qml.numpy.ndarray): measurement_processes = tuple(self.tape.measurements) elif not isinstance(self._qfunc_output, Sequence): diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py index 4fe3d4946d2..f314a80f64b 100644 --- a/tests/interfaces/test_jax_jit.py +++ b/tests/interfaces/test_jax_jit.py @@ -895,6 +895,20 @@ def cost(x, y, device, interface, ek): assert jax.numpy.allclose(r, e, atol=1e-7) +def test_jit_allcounts(): + """Test jitting with counts with all_outcomes == True.""" + + tape = qml.tape.QuantumScript([], [qml.counts(wires=(0, 1), all_outcomes=True)], shots=50) + device = qml.device("default.qubit") + + res = jax.jit(qml.execute, static_argnums=(1, 2))((tape,), device, qml.gradients.param_shift)[0] + + assert set(res.keys()) == {"00", "01", "10", "11"} + assert qml.math.allclose(res["00"], 50) + for val in ["01", "10", "11"]: + assert qml.math.allclose(res[val], 0) + + def test_diff_method_None_jit(): """Test that jitted execution works when `gradient_fn=None`.""" diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py index 20c789fe267..3019da7ddba 100644 --- a/tests/measurements/test_measurements.py +++ b/tests/measurements/test_measurements.py @@ -687,3 +687,50 @@ def circuit(): return CountTapesMP(wires=[0]) assert circuit() == 1 + + +class TestMeasurementProcess: + """Tests for the shape and numeric type of a measurement process""" + + measurements_no_shots = [ + (qml.expval(qml.PauliZ(0)), ()), + (qml.var(qml.PauliZ(0)), ()), + (qml.probs(wires=[0, 1]), (4,)), + (qml.state(), (8,)), + (qml.density_matrix(wires=[0, 1]), (4, 4)), + (qml.mutual_info(wires0=[0], wires1=[1]), ()), + (qml.vn_entropy(wires=[0, 1]), ()), + ] + + measurements_finite_shots = [ + (qml.expval(qml.PauliZ(0)), ()), + (qml.var(qml.PauliZ(0)), ()), + (qml.probs(wires=[0, 1]), (4,)), + (qml.state(), (8,)), + (qml.density_matrix(wires=[0, 1]), (4, 4)), + (qml.sample(qml.PauliZ(0)), (10,)), + (qml.sample(), (10, 3)), + (qml.mutual_info(wires0=0, wires1=1), ()), + (qml.vn_entropy(wires=[0, 1]), ()), + ] + + @pytest.mark.parametrize("measurement, expected_shape", measurements_no_shots) + def test_output_shapes_no_shots(self, measurement, expected_shape): + """Test that the output shape of the measurement process is expected + when shots=None""" + + assert measurement.shape(shots=None, num_device_wires=3) == expected_shape + + @pytest.mark.parametrize("measurement, expected_shape", measurements_finite_shots) + def test_output_shapes_finite_shots(self, measurement, expected_shape): + """Test that the output shape of the measurement process is expected + when shots is finite""" + assert measurement.shape(shots=10, num_device_wires=3) == expected_shape + + def test_undefined_shape_error(self): + """Test that an error is raised for a measurement with an undefined shape""" + measurement = qml.counts(wires=[0, 1]) + msg = "The shape of the measurement CountsMP is not defined" + + with pytest.raises(qml.QuantumFunctionError, match=msg): + measurement.shape(shots=None, num_device_wires=2) diff --git a/tests/tape/test_qscript.py b/tests/tape/test_qscript.py index f72b9f98651..b4ca52bf9a2 100644 --- a/tests/tape/test_qscript.py +++ b/tests/tape/test_qscript.py @@ -991,83 +991,6 @@ def test_diagonalizing_gates_not_queued(self): ] -class TestMeasurementProcess: - """Tests for the shape and numeric type of a measurement process""" - - measurements_no_shots = [ - (qml.expval(qml.PauliZ(0)), ()), - (qml.var(qml.PauliZ(0)), ()), - (qml.probs(wires=[0, 1]), (4,)), - (qml.state(), (8,)), - (qml.density_matrix(wires=[0, 1]), (4, 4)), - (qml.mutual_info(wires0=[0], wires1=[1]), ()), - (qml.vn_entropy(wires=[0, 1]), ()), - ] - - measurements_finite_shots = [ - (qml.expval(qml.PauliZ(0)), ()), - (qml.var(qml.PauliZ(0)), ()), - (qml.probs(wires=[0, 1]), (4,)), - (qml.state(), (8,)), - (qml.density_matrix(wires=[0, 1]), (4, 4)), - (qml.sample(qml.PauliZ(0)), (10,)), - (qml.sample(), (10, 3)), - (qml.mutual_info(wires0=0, wires1=1), ()), - (qml.vn_entropy(wires=[0, 1]), ()), - ] - - measurements_shot_vector = [ - (qml.expval(qml.PauliZ(0)), ((), (), ())), - (qml.var(qml.PauliZ(0)), ((), (), ())), - (qml.probs(wires=[0, 1]), ((4,), (4,), (4,))), - (qml.state(), ((8,), (8,), (8,))), - (qml.density_matrix(wires=[0, 1]), ((4, 4), (4, 4), (4, 4))), - (qml.sample(qml.PauliZ(0)), ((10,), (20,), (30,))), - (qml.sample(), ((10, 3), (20, 3), (30, 3))), - (qml.mutual_info(wires0=0, wires1=1), ((), (), ())), - (qml.vn_entropy(wires=[0, 1]), ((), (), ())), - ] - - @pytest.mark.parametrize("measurement, expected_shape", measurements_no_shots) - def test_output_shapes_no_shots(self, measurement, expected_shape): - """Test that the output shape of the measurement process is expected - when shots=None""" - num_wires = 3 - dev = qml.device("default.qubit", wires=num_wires, shots=None) - - assert measurement.shape(dev, Shots(None)) == expected_shape - - @pytest.mark.parametrize("measurement, expected_shape", measurements_finite_shots) - def test_output_shapes_finite_shots(self, measurement, expected_shape): - """Test that the output shape of the measurement process is expected - when shots is finite""" - num_wires = 3 - num_shots = 10 - dev = qml.device("default.qubit", wires=num_wires, shots=num_shots) - - assert measurement.shape(dev, Shots(num_shots)) == expected_shape - - @pytest.mark.parametrize("measurement, expected_shape", measurements_shot_vector) - def test_output_shapes_shot_vector(self, measurement, expected_shape): - """Test that the output shape of the measurement process is expected - when shots is a vector""" - num_wires = 3 - shot_vector = [10, 20, 30] - dev = qml.device("default.qubit", wires=num_wires, shots=shot_vector) - - assert measurement.shape(dev, Shots(shot_vector)) == expected_shape - - def test_undefined_shape_error(self): - """Test that an error is raised for a measurement with an undefined shape""" - measurement = qml.counts(wires=[0, 1]) - dev = qml.device("default.qubit", wires=2) - shots = Shots(None) - msg = "The shape of the measurement CountsMP is not defined" - - with pytest.raises(qml.QuantumFunctionError, match=msg): - measurement.shape(dev, shots) - - class TestOutputShape: """Tests for determining the tape output shape of tapes.""" @@ -1128,51 +1051,6 @@ def test_output_shapes_single_qnode_check(self, measurement, expected_shape, sho assert qs.shape(dev) == res_shape - def test_output_shapes_single_qnode_check_cutoff(self): - """Test that the tape output shape is correct when computing - probabilities with a dummy device that defines a cutoff value.""" - - class CustomDevice(qml.QubitDevice): - """A dummy device that has a cutoff value specified and returns - analytic probabilities in a fashion similar to the - strawberryfields.fock device. - - Note: this device definition is used as PennyLane-SF is not a - dependency of PennyLane core and there are no CV device in - PennyLane core using a cutoff value. - """ - - name = "Device with cutoff" - short_name = "dummy.device" - pennylane_requires = "0.1.0" - version = "0.0.1" - author = "CV quantum" - - operations = {} - observables = {"Identity"} - - def __init__(self, shots=None, wires=None, cutoff=None): - super().__init__(wires=wires, shots=shots) - self.cutoff = cutoff - - def apply(self, operations, **kwargs): - pass - - def analytic_probability(self, wires=None): - if wires is None: - wires = self.wires - return np.zeros(self.cutoff ** len(wires)) - - dev = CustomDevice(wires=2, cutoff=13) - - # If PennyLane-SF is installed, the following can be checked e.g., locally: - # dev = qml.device("strawberryfields.fock", wires=2, cutoff_dim=13) - - qs = QuantumScript(measurements=[qml.probs(wires=[0])]) - - res_shape = qml.execute([qs], dev, gradient_fn=None)[0] - assert qs.shape(dev) == res_shape.shape - @pytest.mark.autograd @pytest.mark.parametrize("measurements, expected", multi_measurements) @pytest.mark.parametrize("shots", [None, 1, 10]) @@ -1374,16 +1252,10 @@ def test_raises_broadcasting_shot_vector(self): broadcasting along with a device with a shot vector raises an error.""" dev = qml.device("default.qubit", wires=3, shots=(1, 2, 3)) - with qml.queuing.AnnotatedQueue() as q: - qml.RY(np.array([0.1, 0.2]), wires=0) - qml.RX(np.array([0.3, 0.4]), wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q, shots=(1, 2, 3)) - msg = "Parameter broadcasting when using a shot vector is not supported yet" + y = np.array([0.1, 0.2]) + tape = qml.tape.QuantumScript([qml.RY(y, 0)], [qml.expval(qml.Z(0))], shots=(1, 2, 3)) - with pytest.raises(NotImplementedError, match=msg): - tape.shape(dev) + assert tape.shape(dev) == ((2,), (2,), (2,)) class TestNumericType: diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py index b2c70265b91..68866be3284 100644 --- a/tests/tape/test_tape.py +++ b/tests/tape/test_tape.py @@ -1893,54 +1893,6 @@ def test_output_shapes_single_qnode_check(self, measurement, _, shots): res_shape = res_shape if res_shape != tuple() else () assert tape.shape(dev) == res_shape - def test_output_shapes_single_qnode_check_cutoff(self): - """Test that the tape output shape is correct when computing - probabilities with a dummy device that defines a cutoff value.""" - - class CustomDevice(qml.QubitDevice): - """A dummy device that has a cutoff value specified and returns - analytic probabilities in a fashion similar to the - strawberryfields.fock device. - - Note: this device definition is used as PennyLane-SF is not a - dependency of PennyLane core and there are no CV device in - PennyLane core using a cutoff value. - """ - - name = "Device with cutoff" - short_name = "dummy.device" - pennylane_requires = "0.1.0" - version = "0.0.1" - author = "CV quantum" - - operations = {} - observables = {"Identity"} - - def __init__(self, shots=None, wires=None, cutoff=None): - super().__init__(wires=wires, shots=shots) - self.cutoff = cutoff - - def apply(self, operations, **kwargs): - pass - - def analytic_probability(self, wires=None): - if wires is None: - wires = self.wires - return np.zeros(self.cutoff ** len(wires)) - - dev = CustomDevice(wires=2, cutoff=13) - - # If PennyLane-SF is installed, the following can be checked e.g., locally: - # dev = qml.device("strawberryfields.fock", wires=2, cutoff_dim=13) - - with qml.queuing.AnnotatedQueue() as q: - qml.probs(wires=[0]) - - tape = qml.tape.QuantumScript.from_queue(q) - - res_shape = qml.execute([tape], dev, gradient_fn=qml.gradients.param_shift_cv)[0] - assert tape.shape(dev) == res_shape.shape - @pytest.mark.autograd @pytest.mark.parametrize("measurements, expected", multi_measurements) @pytest.mark.parametrize("shots", [None, 1, 10]) diff --git a/tests/test_qnode.py b/tests/test_qnode.py index 9f452bfe853..f1bb08ecb38 100644 --- a/tests/test_qnode.py +++ b/tests/test_qnode.py @@ -708,7 +708,7 @@ def circuit1(param): jitted_qnode1 = jax.jit(qn) with pytest.raises( - qml.QuantumFunctionError, match="Can't JIT a quantum function that returns counts." + NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts." ): jitted_qnode1(0.123) @@ -723,7 +723,7 @@ def circuit2(param): jitted_qnode2 = jax.jit(circuit2) with pytest.raises( - qml.QuantumFunctionError, match="Can't JIT a quantum function that returns counts." + NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts." ): jitted_qnode2(0.123) diff --git a/tests/test_qnode_legacy.py b/tests/test_qnode_legacy.py index 7ea5c331596..4f166b3e1b9 100644 --- a/tests/test_qnode_legacy.py +++ b/tests/test_qnode_legacy.py @@ -891,7 +891,7 @@ def circuit1(param): jitted_qnode1 = jax.jit(qn) with pytest.raises( - qml.QuantumFunctionError, match="Can't JIT a quantum function that returns counts." + NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts." ): jitted_qnode1(0.123) @@ -906,7 +906,7 @@ def circuit2(param): jitted_qnode2 = jax.jit(circuit2) with pytest.raises( - qml.QuantumFunctionError, match="Can't JIT a quantum function that returns counts." + NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts." ): jitted_qnode2(0.123) From b30e4cc02dd0352b8e76d852472567be64a6b1c6 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Fri, 16 Aug 2024 10:43:02 -0400 Subject: [PATCH 2/5] test improvement --- tests/interfaces/test_jax_jit.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py index f314a80f64b..3bcac7055a1 100644 --- a/tests/interfaces/test_jax_jit.py +++ b/tests/interfaces/test_jax_jit.py @@ -898,7 +898,9 @@ def cost(x, y, device, interface, ek): def test_jit_allcounts(): """Test jitting with counts with all_outcomes == True.""" - tape = qml.tape.QuantumScript([], [qml.counts(wires=(0, 1), all_outcomes=True)], shots=50) + tape = qml.tape.QuantumScript( + [qml.RX(0, 0)], [qml.counts(wires=(0, 1), all_outcomes=True)], shots=50 + ) device = qml.device("default.qubit") res = jax.jit(qml.execute, static_argnums=(1, 2))((tape,), device, qml.gradients.param_shift)[0] From 045eee7ec68693e6bd5b8921fefcae92b1f4c11d Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Fri, 16 Aug 2024 10:52:08 -0400 Subject: [PATCH 3/5] Apply suggestions from code review --- doc/releases/changelog-dev.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index dbf88be8fef..a9a2662be89 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -30,6 +30,7 @@ * Counts measurements with `all_outcomes=True` can now be used with jax jitting. Measurements broadcasted across all available wires (`qml.probs()`) can now be used with jit and devices that allow variable numbers of wires (`qml.device('default.qubit')`). + [(#6108)](https://github.com/PennyLaneAI/pennylane/pull/6108/) * Mid-circuit measurements can now be captured with `qml.capture` enabled. [(#6015)](https://github.com/PennyLaneAI/pennylane/pull/6015) @@ -238,9 +239,11 @@ * `MeasurementProcess.shape(shots: Shots, device:Device)` is now `MeasurementProcess.shape(shots: Optional[int], num_device_wires:int = 0)`. This has been done to allow jitting when a measurement is broadcasted across all available wires, but the device does not specify wires. + [(#6108)](https://github.com/PennyLaneAI/pennylane/pull/6108/) * If the shape of a probability measurement is affected by a `Device.cutoff` property, it will no longer work with jitting. + [(#6108)](https://github.com/PennyLaneAI/pennylane/pull/6108/) * `GlobalPhase` is considered non-differentiable with tape transforms. As a consequence, `qml.gradients.finite_diff` and `qml.gradients.spsa_grad` no longer From ec63f2ed95f5c7846f4807f8ccc16605963cbaa2 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Fri, 16 Aug 2024 14:40:18 -0400 Subject: [PATCH 4/5] fixing tests --- .../legacy/test_classical_shadow_legacy.py | 29 +------- .../measurements/legacy/test_expval_legacy.py | 26 +------ .../legacy/test_mutual_info_legacy.py | 9 --- .../measurements/legacy/test_probs_legacy.py | 22 +----- .../legacy/test_purity_measurement_legacy.py | 9 --- .../measurements/legacy/test_sample_legacy.py | 72 +++---------------- .../measurements/legacy/test_state_legacy.py | 36 +--------- tests/measurements/legacy/test_var_legacy.py | 25 +------ .../legacy/test_vn_entropy_legacy.py | 9 --- tests/measurements/test_classical_shadow.py | 17 ++--- tests/measurements/test_expval.py | 18 +---- tests/measurements/test_mutual_info.py | 7 +- tests/measurements/test_probs.py | 23 ++---- tests/measurements/test_purity_measurement.py | 7 +- tests/measurements/test_sample.py | 47 ++++-------- tests/measurements/test_state.py | 34 +-------- tests/measurements/test_var.py | 18 +---- tests/measurements/test_vn_entropy.py | 7 +- 18 files changed, 56 insertions(+), 359 deletions(-) diff --git a/tests/measurements/legacy/test_classical_shadow_legacy.py b/tests/measurements/legacy/test_classical_shadow_legacy.py index e049247e717..14b7485254c 100644 --- a/tests/measurements/legacy/test_classical_shadow_legacy.py +++ b/tests/measurements/legacy/test_classical_shadow_legacy.py @@ -18,7 +18,6 @@ import pennylane as qml from pennylane import numpy as np -from pennylane.measurements import Shots # pylint: disable=dangerous-default-value, too-many-arguments @@ -120,20 +119,6 @@ class TestClassicalShadow: shots_list = [1, 100] seed_recipes_list = [None, 74] # random seed - @pytest.mark.parametrize("shots", shots_list) - @pytest.mark.parametrize("seed", seed_recipes_list) - def test_measurement_process_shape(self, wires, shots, seed): - """Test that the shape of the MeasurementProcess instance is correct""" - dev = qml.device("default.qubit.legacy", wires=wires, shots=shots) - shots_obj = Shots(shots) - res = qml.classical_shadow(wires=range(wires), seed=seed) - assert res.shape(dev, shots_obj) == (2, shots, wires) - - # test an error is raised when device is None - msg = "Shots must be specified to obtain the shape of a classical shadow measurement" - with pytest.raises(qml.measurements.MeasurementShapeError, match=msg): - res.shape(dev, Shots(None)) - def test_shape_matches(self, wires): """Test that the shape of the MeasurementProcess matches the shape of the tape execution""" @@ -143,9 +128,7 @@ def test_shape_matches(self, wires): circuit.construct((), {}) res = qml.execute([circuit.tape], circuit.device, None)[0] - expected_shape = qml.classical_shadow(wires=range(wires)).shape( - circuit.device, Shots(shots) - ) + expected_shape = qml.classical_shadow(wires=range(wires)).shape(shots, wires) assert res.shape == expected_shape @@ -286,14 +269,6 @@ def circuit(obs, k=1): @pytest.mark.autograd class TestExpvalMeasurement: - @pytest.mark.parametrize("wires", [1, 2]) - @pytest.mark.parametrize("shots", [1, 10]) - def test_measurement_process_shape(self, wires, shots): - """Test that the shape of the MeasurementProcess instance is correct""" - dev = qml.device("default.qubit.legacy", wires=wires, shots=shots) - H = qml.PauliZ(0) - res = qml.shadow_expval(H) - assert len(res.shape(dev, Shots(shots))) == 0 def test_shape_matches(self): """Test that the shape of the MeasurementProcess matches the shape @@ -306,7 +281,7 @@ def test_shape_matches(self): circuit.construct((H,), {}) res = qml.execute([circuit.tape], circuit.device, None)[0] - expected_shape = qml.shadow_expval(H).shape(circuit.device, Shots(shots)) + expected_shape = qml.shadow_expval(H).shape(shots, len(circuit.device.wires)) assert res.shape == expected_shape diff --git a/tests/measurements/legacy/test_expval_legacy.py b/tests/measurements/legacy/test_expval_legacy.py index da59c0aaca8..5be4ba90d59 100644 --- a/tests/measurements/legacy/test_expval_legacy.py +++ b/tests/measurements/legacy/test_expval_legacy.py @@ -17,7 +17,7 @@ import pennylane as qml from pennylane.devices.qubit.measure import flatten_state -from pennylane.measurements import Expectation, Shots +from pennylane.measurements import Expectation # TODO: Remove this when new CustomMP are the default @@ -169,30 +169,6 @@ def expected_circuit(phi): if device_name != "default.mixed": custom_measurement_process(new_dev, spy) - @pytest.mark.parametrize( - "obs", - [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)], - ) - def test_shape(self, obs): - """Test that the shape is correct.""" - dev = qml.device("default.qubit.legacy", wires=1) - - res = qml.expval(obs) - # pylint: disable=use-implicit-booleaness-not-comparison - assert res.shape(dev, Shots(None)) == () - assert res.shape(dev, Shots(100)) == () - - @pytest.mark.parametrize( - "obs", - [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)], - ) - def test_shape_shot_vector(self, obs): - """Test that the shape is correct with the shot vector too.""" - res = qml.expval(obs) - shot_vector = (1, 2, 3) - dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector) - assert res.shape(dev, Shots(shot_vector)) == ((), (), ()) - @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])]) @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]]) def test_projector_expval(self, state, shots, mocker): diff --git a/tests/measurements/legacy/test_mutual_info_legacy.py b/tests/measurements/legacy/test_mutual_info_legacy.py index 023c79205e7..fee7b9f071b 100644 --- a/tests/measurements/legacy/test_mutual_info_legacy.py +++ b/tests/measurements/legacy/test_mutual_info_legacy.py @@ -16,18 +16,9 @@ import pytest import pennylane as qml -from pennylane.measurements import Shots from pennylane.workflow import INTERFACE_MAP -@pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ([1, 10], ((), ()))]) -def test_shape(shots, shape): - """Test that the shape is correct.""" - dev = qml.device("default.qubit.legacy", wires=3, shots=shots) - res = qml.mutual_info(wires0=[0], wires1=[1]) - assert res.shape(dev, Shots(shots)) == shape - - class TestIntegration: """Tests for the mutual information functions""" diff --git a/tests/measurements/legacy/test_probs_legacy.py b/tests/measurements/legacy/test_probs_legacy.py index 2fe3243bab8..d2db413d396 100644 --- a/tests/measurements/legacy/test_probs_legacy.py +++ b/tests/measurements/legacy/test_probs_legacy.py @@ -18,7 +18,7 @@ import pennylane as qml from pennylane import numpy as pnp from pennylane.devices.qubit.measure import flatten_state -from pennylane.measurements import ProbabilityMP, Shots +from pennylane.measurements import ProbabilityMP # TODO: Remove this when new CustomMP are the default @@ -79,26 +79,6 @@ def circuit(): assert isinstance(circuit.tape[0], ProbabilityMP) - @pytest.mark.parametrize("wires", [[0], [2, 1], ["a", "c", 3]]) - @pytest.mark.parametrize("shots", [None, 10]) - def test_shape(self, wires, shots): - """Test that the shape is correct.""" - dev = qml.device("default.qubit.legacy", wires=3, shots=shots) - res = qml.probs(wires=wires) - assert res.shape(dev, Shots(shots)) == (2 ** len(wires),) - - @pytest.mark.parametrize("wires", [[0], [2, 1], ["a", "c", 3]]) - def test_shape_shot_vector(self, wires): - """Test that the shape is correct with the shot vector too.""" - res = qml.probs(wires=wires) - shot_vector = (1, 2, 3) - dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector) - assert res.shape(dev, Shots(shot_vector)) == ( - (2 ** len(wires),), - (2 ** len(wires),), - (2 ** len(wires),), - ) - @pytest.mark.parametrize("shots", [None, 100]) def test_probs_no_arguments(self, shots): """Test that using ``qml.probs`` with no arguments returns the probabilities of all wires.""" diff --git a/tests/measurements/legacy/test_purity_measurement_legacy.py b/tests/measurements/legacy/test_purity_measurement_legacy.py index 752b07b5f7b..4ff7a850823 100644 --- a/tests/measurements/legacy/test_purity_measurement_legacy.py +++ b/tests/measurements/legacy/test_purity_measurement_legacy.py @@ -17,7 +17,6 @@ import pytest import pennylane as qml -from pennylane.measurements import Shots # pylint: disable=too-many-arguments @@ -53,14 +52,6 @@ def expected_purity_grad_ising_xx(param): return grad_expected_purity -@pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))]) -def test_shape_new(shots, shape): - """Test the ``shape_new`` method.""" - meas = qml.purity(wires=0) - dev = qml.device("default.qubit.legacy", wires=1, shots=shots) - assert meas.shape(dev, Shots(shots)) == shape - - class TestPurityIntegration: """Test the purity meausrement with qnodes and devices.""" diff --git a/tests/measurements/legacy/test_sample_legacy.py b/tests/measurements/legacy/test_sample_legacy.py index b1bc69bb70d..9f836e7764c 100644 --- a/tests/measurements/legacy/test_sample_legacy.py +++ b/tests/measurements/legacy/test_sample_legacy.py @@ -16,7 +16,7 @@ import pytest import pennylane as qml -from pennylane.measurements import MeasurementShapeError, MeasurementValue, Sample, Shots +from pennylane.measurements import MeasurementValue, Sample from pennylane.operation import Operator # pylint: disable=protected-access, no-member @@ -62,10 +62,10 @@ def circuit(): output = circuit() assert len(output) == 2 - assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == ( + assert circuit._qfunc_output[0].shape(n_sample, 2) == ( (n_sample,) if not n_sample == 1 else () ) - assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == ( + assert circuit._qfunc_output[1].shape(n_sample, 2) == ( (n_sample,) if not n_sample == 1 else () ) @@ -89,7 +89,7 @@ def circuit(): assert len(result) == 3 assert np.array_equal(result[0].shape, (n_sample,)) - assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,) + assert circuit._qfunc_output[0].shape(n_sample, 3) == (n_sample,) assert isinstance(result[1], np.ndarray) assert isinstance(result[2], np.ndarray) @@ -198,7 +198,7 @@ def circuit(): assert isinstance(result, np.ndarray) assert np.array_equal(result.shape, (n_sample,)) - assert circuit._qfunc_output.shape(dev, Shots(n_sample)) == (n_sample,) + assert circuit._qfunc_output.shape(n_sample, 1) == (n_sample,) custom_measurement_process(dev, spy) @@ -216,9 +216,9 @@ def circuit(): result = circuit() - assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,) - assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == (n_sample,) - assert circuit._qfunc_output[2].shape(dev, Shots(n_sample)) == (n_sample,) + assert circuit._qfunc_output[0].shape(n_sample, 3) == (n_sample,) + assert circuit._qfunc_output[1].shape(n_sample, 3) == (n_sample,) + assert circuit._qfunc_output[2].shape(n_sample, 3) == (n_sample,) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) @@ -353,58 +353,6 @@ def circuit(): custom_measurement_process(dev, spy) - def test_shape_no_shots_error(self): - """Test that the appropriate error is raised with no shots are specified""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) - shots = Shots(None) - mp = qml.sample() - - with pytest.raises( - MeasurementShapeError, match="Shots are required to obtain the shape of the measurement" - ): - _ = mp.shape(dev, shots) - - @pytest.mark.parametrize( - "obs", - [ - None, - qml.PauliZ(0), - qml.Hermitian(np.diag([1, 2]), 0), - qml.Hermitian(np.diag([1.0, 2.0]), 0), - ], - ) - def test_shape(self, obs): - """Test that the shape is correct.""" - shots = 10 - dev = qml.device("default.qubit.legacy", wires=3, shots=shots) - res = qml.sample(obs) if obs is not None else qml.sample() - expected = (shots,) if obs is not None else (shots, 3) - assert res.shape(dev, Shots(shots)) == expected - - @pytest.mark.parametrize("n_samples", (1, 10)) - def test_shape_wires(self, n_samples): - """Test that the shape is correct when wires are provided.""" - dev = qml.device("default.qubit.legacy", wires=3, shots=n_samples) - mp = qml.sample(wires=(0, 1)) - assert mp.shape(dev, Shots(n_samples)) == (n_samples, 2) if n_samples != 1 else (2,) - - @pytest.mark.parametrize( - "obs", - [ - None, - qml.PauliZ(0), - qml.Hermitian(np.diag([1, 2]), 0), - qml.Hermitian(np.diag([1.0, 2.0]), 0), - ], - ) - def test_shape_shot_vector(self, obs): - """Test that the shape is correct with the shot vector too.""" - shot_vector = (1, 2, 3) - dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector) - res = qml.sample(obs) if obs is not None else qml.sample() - expected = ((), (2,), (3,)) if obs is not None else ((3,), (2, 3), (3, 3)) - assert res.shape(dev, Shots(shot_vector)) == expected - def test_shape_shot_vector_obs(self): """Test that the shape is correct with the shot vector and a observable too.""" shot_vec = (2, 2) @@ -459,6 +407,4 @@ def circuit(x): expected = (2,) if samples == 1 else (samples, 2) assert results.shape == expected - assert ( - circuit._qfunc_output.shape(dev, Shots(samples)) == (samples, 2) if samples != 1 else (2,) - ) + assert circuit._qfunc_output.shape(samples, 3) == (samples, 2) if samples != 1 else (2,) diff --git a/tests/measurements/legacy/test_state_legacy.py b/tests/measurements/legacy/test_state_legacy.py index 83e7cba2b71..6e9cadcae13 100644 --- a/tests/measurements/legacy/test_state_legacy.py +++ b/tests/measurements/legacy/test_state_legacy.py @@ -18,7 +18,7 @@ import pennylane as qml from pennylane import numpy as pnp from pennylane.devices import DefaultQubitLegacy -from pennylane.measurements import Shots, State, density_matrix, expval, state +from pennylane.measurements import State, density_matrix, expval, state class TestState: @@ -324,20 +324,6 @@ def func(): assert np.allclose(state_expected, state_val) - @pytest.mark.parametrize("shots", [None, 1, 10]) - def test_shape(self, shots): - """Test that the shape is correct for qml.state.""" - dev = qml.device("default.qubit.legacy", wires=3, shots=shots) - res = qml.state() - assert res.shape(dev, Shots(shots)) == (2**3,) - - @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)]) - def test_shape_shot_vector(self, s_vec): - """Test that the shape is correct for qml.state with the shot vector too.""" - dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec) - res = qml.state() - assert res.shape(dev, Shots(s_vec)) == ((2**3,), (2**3,), (2**3,)) - class TestDensityMatrix: """Tests for the density matrix function""" @@ -373,7 +359,7 @@ def func(): return density_matrix(0) func() - obs = func.qtape.observables + obs = func.qtape.measurements assert len(obs) == 1 assert obs[0].return_type is State @@ -783,21 +769,3 @@ def func(): ), density, ) - - @pytest.mark.parametrize("shots", [None, 1, 10]) - def test_shape(self, shots): - """Test that the shape is correct for qml.density_matrix.""" - dev = qml.device("default.qubit.legacy", wires=3, shots=shots) - res = qml.density_matrix(wires=[0, 1]) - assert res.shape(dev, Shots(shots)) == (2**2, 2**2) - - @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)]) - def test_shape_shot_vector(self, s_vec): - """Test that the shape is correct for qml.density_matrix with the shot vector too.""" - dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec) - res = qml.density_matrix(wires=[0, 1]) - assert res.shape(dev, Shots(s_vec)) == ( - (2**2, 2**2), - (2**2, 2**2), - (2**2, 2**2), - ) diff --git a/tests/measurements/legacy/test_var_legacy.py b/tests/measurements/legacy/test_var_legacy.py index 4472d01dc22..985f75d524a 100644 --- a/tests/measurements/legacy/test_var_legacy.py +++ b/tests/measurements/legacy/test_var_legacy.py @@ -17,7 +17,7 @@ import pennylane as qml from pennylane.devices.qubit.measure import flatten_state -from pennylane.measurements import Shots, Variance +from pennylane.measurements import Variance # TODO: Remove this when new CustomMP are the default @@ -166,29 +166,6 @@ def expected_circuit(phi): if device_name != "default.mixed": custom_measurement_process(new_dev, spy) - @pytest.mark.parametrize( - "obs", - [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)], - ) - def test_shape(self, obs): - """Test that the shape is correct.""" - dev = qml.device("default.qubit.legacy", wires=1) - res = qml.var(obs) - # pylint: disable=use-implicit-booleaness-not-comparison - assert res.shape(dev, Shots(None)) == () - assert res.shape(dev, Shots(100)) == () - - @pytest.mark.parametrize( - "obs", - [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)], - ) - def test_shape_shot_vector(self, obs): - """Test that the shape is correct with the shot vector too.""" - res = qml.var(obs) - shot_vector = (1, 2, 3) - dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector) - assert res.shape(dev, Shots(shot_vector)) == ((), (), ()) - @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])]) @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]]) def test_projector_var(self, state, shots, mocker): diff --git a/tests/measurements/legacy/test_vn_entropy_legacy.py b/tests/measurements/legacy/test_vn_entropy_legacy.py index 9a4dbf2f7d9..6a9f53b3b85 100644 --- a/tests/measurements/legacy/test_vn_entropy_legacy.py +++ b/tests/measurements/legacy/test_vn_entropy_legacy.py @@ -16,7 +16,6 @@ import pytest import pennylane as qml -from pennylane.measurements import Shots from pennylane.measurements.vn_entropy import VnEntropyMP from pennylane.workflow import INTERFACE_MAP @@ -101,14 +100,6 @@ def circuit(): assert isinstance(circuit.tape[0], VnEntropyMP) - @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))]) - def test_shape(self, shots, shape): - """Test the ``shape`` method.""" - meas = qml.vn_entropy(wires=0) - dev = qml.device("default.qubit.legacy", wires=1, shots=shots) - - assert meas.shape(dev, Shots(shots)) == shape - class TestIntegration: """Integration tests for the vn_entropy measurement function.""" diff --git a/tests/measurements/test_classical_shadow.py b/tests/measurements/test_classical_shadow.py index 04d9f8e80af..a900a131337 100644 --- a/tests/measurements/test_classical_shadow.py +++ b/tests/measurements/test_classical_shadow.py @@ -20,7 +20,7 @@ import pennylane as qml from pennylane import numpy as np -from pennylane.measurements import ClassicalShadowMP, Shots +from pennylane.measurements import ClassicalShadowMP from pennylane.measurements.classical_shadow import ShadowExpvalMP # pylint: disable=dangerous-default-value, too-many-arguments @@ -239,15 +239,13 @@ def test_measurement_process_numeric_type(self, wires, seed): @pytest.mark.parametrize("seed", seed_recipes_list) def test_measurement_process_shape(self, wires, shots, seed): """Test that the shape of the MeasurementProcess instance is correct""" - dev = qml.device("default.qubit", wires=wires, shots=shots) - shots_obj = Shots(shots) res = qml.classical_shadow(wires=range(wires), seed=seed) - assert res.shape(dev, shots_obj) == (2, shots, wires) + assert res.shape(shots, wires) == (2, shots, wires) # test an error is raised when device is None msg = "Shots must be specified to obtain the shape of a classical shadow measurement" with pytest.raises(qml.measurements.MeasurementShapeError, match=msg): - res.shape(dev, Shots(None)) + res.shape(None, wires) def test_shape_matches(self, wires): """Test that the shape of the MeasurementProcess matches the shape @@ -258,9 +256,7 @@ def test_shape_matches(self, wires): circuit.construct((), {}) res = qml.execute([circuit.tape], circuit.device, None)[0] - expected_shape = qml.classical_shadow(wires=range(wires)).shape( - circuit.device, Shots(shots) - ) + expected_shape = qml.classical_shadow(wires=range(wires)).shape(shots, wires) assert res.shape == expected_shape @@ -422,10 +418,9 @@ def test_measurement_process_numeric_type(self): @pytest.mark.parametrize("shots", [1, 10]) def test_measurement_process_shape(self, wires, shots): """Test that the shape of the MeasurementProcess instance is correct""" - dev = qml.device("default.qubit", wires=wires, shots=shots) H = qml.PauliZ(0) res = qml.shadow_expval(H) - assert len(res.shape(dev, Shots(shots))) == 0 + assert len(res.shape(shots, wires)) == 0 def test_shape_matches(self): """Test that the shape of the MeasurementProcess matches the shape @@ -438,7 +433,7 @@ def test_shape_matches(self): circuit.construct((H,), {}) res = qml.execute([circuit.tape], circuit.device, None)[0] - expected_shape = qml.shadow_expval(H).shape(circuit.device, Shots(shots)) + expected_shape = qml.shadow_expval(H).shape(shots, wires) assert res.shape == expected_shape diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py index 14dcbb4ba83..b1a5e717832 100644 --- a/tests/measurements/test_expval.py +++ b/tests/measurements/test_expval.py @@ -18,7 +18,7 @@ import pytest import pennylane as qml -from pennylane.measurements import Expectation, Shots +from pennylane.measurements import Expectation from pennylane.measurements.expval import ExpectationMP @@ -180,23 +180,11 @@ def test_numeric_type(self, obs): ) def test_shape(self, obs): """Test that the shape is correct.""" - dev = qml.device("default.qubit", wires=1) res = qml.expval(obs) # pylint: disable=use-implicit-booleaness-not-comparison - assert res.shape(dev, Shots(None)) == () - assert res.shape(dev, Shots(100)) == () - - @pytest.mark.parametrize( - "obs", - [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)], - ) - def test_shape_shot_vector(self, obs): - """Test that the shape is correct with the shot vector too.""" - res = qml.expval(obs) - shot_vector = (1, 2, 3) - dev = qml.device("default.qubit", wires=3, shots=shot_vector) - assert res.shape(dev, Shots(shot_vector)) == ((), (), ()) + assert res.shape(None, 1) == () + assert res.shape(100, 1) == () @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])]) @pytest.mark.parametrize("shots", [None, 1000, [1000, 1111]]) diff --git a/tests/measurements/test_mutual_info.py b/tests/measurements/test_mutual_info.py index c457afd6b65..a47284f744b 100644 --- a/tests/measurements/test_mutual_info.py +++ b/tests/measurements/test_mutual_info.py @@ -18,7 +18,7 @@ import pytest import pennylane as qml -from pennylane.measurements import MutualInfo, Shots +from pennylane.measurements import MutualInfo from pennylane.measurements.mutual_info import MutualInfoMP from pennylane.wires import Wires @@ -35,12 +35,11 @@ def test_queue(self): assert q.queue[0] is m assert isinstance(q.queue[0], MutualInfoMP) - @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ([1, 10], ((), ()))]) + @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ())]) def test_shape(self, shots, shape): """Test that the shape is correct.""" - dev = qml.device("default.qubit", wires=3, shots=shots) res = qml.mutual_info(wires0=[0], wires1=[1]) - assert res.shape(dev, Shots(shots)) == shape + assert res.shape(shots, 3) == shape def test_properties(self): """Test that the properties are correct.""" diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py index 83e76fca5b3..48d40d5d548 100644 --- a/tests/measurements/test_probs.py +++ b/tests/measurements/test_probs.py @@ -19,7 +19,7 @@ import pennylane as qml from pennylane import numpy as pnp -from pennylane.measurements import MeasurementProcess, Probability, ProbabilityMP, Shots +from pennylane.measurements import MeasurementProcess, Probability, ProbabilityMP from pennylane.queuing import AnnotatedQueue @@ -63,31 +63,16 @@ def test_numeric_type(self): @pytest.mark.parametrize("shots", [None, 10]) def test_shape(self, wires, shots): """Test that the shape is correct.""" - dev = qml.device("default.qubit", wires=3, shots=shots) res = qml.probs(wires=wires) - assert res.shape(dev, Shots(shots)) == (2 ** len(wires),) + assert res.shape(shots, 3) == (2 ** len(wires),) def test_shape_empty_wires(self): """Test that shape works when probs is broadcasted onto all available wires.""" - dev = qml.device("default.qubit", wires=(1, 2, 3)) res = qml.probs() - assert res.shape(dev, Shots(None)) == (8,) + assert res.shape(None, 3) == (8,) - dev2 = qml.device("default.qubit") res = qml.probs() - assert res.shape(dev2, Shots(None)) == (1,) - - @pytest.mark.parametrize("wires", [[0], [2, 1], ["a", "c", 3]]) - def test_shape_shot_vector(self, wires): - """Test that the shape is correct with the shot vector too.""" - res = qml.probs(wires=wires) - shot_vector = (1, 2, 3) - dev = qml.device("default.qubit", wires=3, shots=shot_vector) - assert res.shape(dev, Shots(shot_vector)) == ( - (2 ** len(wires),), - (2 ** len(wires),), - (2 ** len(wires),), - ) + assert res.shape(None, 0) == (1,) @pytest.mark.parametrize("wires", [[0], [0, 1], [1, 0, 2]]) def test_annotating_probs(self, wires): diff --git a/tests/measurements/test_purity_measurement.py b/tests/measurements/test_purity_measurement.py index 2ed4a9ea50f..81905102336 100644 --- a/tests/measurements/test_purity_measurement.py +++ b/tests/measurements/test_purity_measurement.py @@ -17,7 +17,7 @@ import pytest import pennylane as qml -from pennylane.measurements import PurityMP, Shots +from pennylane.measurements import PurityMP # pylint: disable=too-many-arguments @@ -66,12 +66,11 @@ def test_numeric_type(self): m = PurityMP(wires=qml.wires.Wires(0)) assert m.numeric_type is float - @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))]) + @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ())]) def test_shape_new(self, shots, shape): """Test the ``shape_new`` method.""" meas = qml.purity(wires=0) - dev = qml.device("default.qubit", wires=1, shots=shots) - assert meas.shape(dev, Shots(shots)) == shape + assert meas.shape(shots, 1) == shape class TestPurityIntegration: diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py index c54b78f620f..fb181123f87 100644 --- a/tests/measurements/test_sample.py +++ b/tests/measurements/test_sample.py @@ -16,7 +16,7 @@ import pytest import pennylane as qml -from pennylane.measurements import MeasurementShapeError, Sample, Shots +from pennylane.measurements import MeasurementShapeError, Sample from pennylane.operation import EigvalsUndefinedError, Operator # pylint: disable=protected-access, no-member, too-many-public-methods @@ -39,10 +39,10 @@ def circuit(): output = circuit() assert len(output) == 2 - assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == ( + assert circuit._qfunc_output[0].shape(shots=n_sample, num_device_wires=2) == ( (n_sample,) if not n_sample == 1 else () ) - assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == ( + assert circuit._qfunc_output[1].shape(shots=n_sample, num_device_wires=2) == ( (n_sample,) if not n_sample == 1 else () ) @@ -62,7 +62,7 @@ def circuit(): assert len(result) == 3 assert np.array_equal(result[0].shape, (n_sample,)) - assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,) + assert circuit._qfunc_output[0].shape(shots=n_sample, num_device_wires=3) == (n_sample,) assert isinstance(result[1], np.float64) assert isinstance(result[2], np.float64) @@ -81,7 +81,7 @@ def circuit(): assert isinstance(result, np.ndarray) assert np.array_equal(result.shape, (n_sample,)) - assert circuit._qfunc_output.shape(dev, Shots(n_sample)) == (n_sample,) + assert circuit._qfunc_output.shape(shots=n_sample, num_device_wires=1) == (n_sample,) def test_multi_wire_sample_regular_shape(self): """Test the return type and shape of sampling multiple wires @@ -96,9 +96,9 @@ def circuit(): result = circuit() - assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,) - assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == (n_sample,) - assert circuit._qfunc_output[2].shape(dev, Shots(n_sample)) == (n_sample,) + assert circuit._qfunc_output[0].shape(shots=n_sample, num_device_wires=3) == (n_sample,) + assert circuit._qfunc_output[1].shape(shots=n_sample, num_device_wires=3) == (n_sample,) + assert circuit._qfunc_output[2].shape(shots=n_sample, num_device_wires=3) == (n_sample,) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) @@ -347,14 +347,12 @@ def test_numeric_type(self, obs): def test_shape_no_shots_error(self): """Test that the appropriate error is raised with no shots are specified""" - dev = qml.device("default.qubit", wires=2, shots=None) - shots = Shots(None) mp = qml.sample() with pytest.raises( MeasurementShapeError, match="Shots are required to obtain the shape of the measurement" ): - _ = mp.shape(dev, shots) + _ = mp.shape(shots=None) @pytest.mark.parametrize( "obs", @@ -368,34 +366,15 @@ def test_shape_no_shots_error(self): def test_shape(self, obs): """Test that the shape is correct.""" shots = 10 - dev = qml.device("default.qubit", wires=3, shots=shots) res = qml.sample(obs) if obs is not None else qml.sample() expected = (shots,) if obs is not None else (shots, 3) - assert res.shape(dev, Shots(shots)) == expected + assert res.shape(10, 3) == expected @pytest.mark.parametrize("n_samples", (1, 10)) def test_shape_wires(self, n_samples): """Test that the shape is correct when wires are provided.""" - dev = qml.device("default.qubit", wires=3, shots=n_samples) mp = qml.sample(wires=(0, 1)) - assert mp.shape(dev, Shots(n_samples)) == (n_samples, 2) if n_samples != 1 else (2,) - - @pytest.mark.parametrize( - "obs", - [ - None, - qml.PauliZ(0), - qml.Hermitian(np.diag([1, 2]), 0), - qml.Hermitian(np.diag([1.0, 2.0]), 0), - ], - ) - def test_shape_shot_vector(self, obs): - """Test that the shape is correct with the shot vector too.""" - shot_vector = (1, 2, 3) - dev = qml.device("default.qubit", wires=3, shots=shot_vector) - res = qml.sample(obs) if obs is not None else qml.sample() - expected = ((), (2,), (3,)) if obs is not None else ((3,), (2, 3), (3, 3)) - assert res.shape(dev, Shots(shot_vector)) == expected + assert mp.shape(n_samples, 3) == (n_samples, 2) if n_samples != 1 else (2,) def test_shape_shot_vector_obs(self): """Test that the shape is correct with the shot vector and a observable too.""" @@ -494,9 +473,7 @@ def circuit(x): expected = (2,) if samples == 1 else (samples, 2) assert results.shape == expected - assert ( - circuit._qfunc_output.shape(dev, Shots(samples)) == (samples, 2) if samples != 1 else (2,) - ) + assert circuit._qfunc_output.shape(samples, 3) == (samples, 2) if samples != 1 else (2,) @pytest.mark.jax diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py index 6dce16479d2..838cad68eab 100644 --- a/tests/measurements/test_state.py +++ b/tests/measurements/test_state.py @@ -20,15 +20,7 @@ from pennylane.devices import DefaultMixed from pennylane.math.matrix_manipulation import _permute_dense_matrix from pennylane.math.quantum import reduce_dm, reduce_statevector -from pennylane.measurements import ( - DensityMatrixMP, - Shots, - State, - StateMP, - density_matrix, - expval, - state, -) +from pennylane.measurements import DensityMatrixMP, State, StateMP, density_matrix, expval, state from pennylane.wires import WireError, Wires @@ -521,16 +513,8 @@ def func(): @pytest.mark.parametrize("shots", [None, 1, 10]) def test_shape(self, shots): """Test that the shape is correct for qml.state.""" - dev = qml.device("default.qubit", wires=3, shots=shots) res = qml.state() - assert res.shape(dev, Shots(shots)) == (2**3,) - - @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)]) - def test_shape_shot_vector(self, s_vec): - """Test that the shape is correct for qml.state with the shot vector too.""" - dev = qml.device("default.qubit", wires=3, shots=s_vec) - res = qml.state() - assert res.shape(dev, Shots(s_vec)) == ((2**3,), (2**3,), (2**3,)) + assert res.shape(shots, 3) == (2**3,) def test_numeric_type(self): """Test that the numeric type of state measurements.""" @@ -1091,17 +1075,5 @@ def func(): @pytest.mark.parametrize("shots", [None, 1, 10]) def test_shape(self, shots): """Test that the shape is correct for qml.density_matrix.""" - dev = qml.device("default.qubit", wires=3, shots=shots) res = qml.density_matrix(wires=[0, 1]) - assert res.shape(dev, Shots(shots)) == (2**2, 2**2) - - @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)]) - def test_shape_shot_vector(self, s_vec): - """Test that the shape is correct for qml.density_matrix with the shot vector too.""" - dev = qml.device("default.qubit", wires=3, shots=s_vec) - res = qml.density_matrix(wires=[0, 1]) - assert res.shape(dev, Shots(s_vec)) == ( - (2**2, 2**2), - (2**2, 2**2), - (2**2, 2**2), - ) + assert res.shape(shots, 3) == (2**2, 2**2) diff --git a/tests/measurements/test_var.py b/tests/measurements/test_var.py index 8d0619a213f..e737a6bbefa 100644 --- a/tests/measurements/test_var.py +++ b/tests/measurements/test_var.py @@ -18,7 +18,7 @@ from flaky import flaky import pennylane as qml -from pennylane.measurements import Shots, Variance, VarianceMP +from pennylane.measurements import Variance, VarianceMP class TestVar: @@ -156,22 +156,10 @@ def test_numeric_type(self, obs): ) def test_shape(self, obs): """Test that the shape is correct.""" - dev = qml.device("default.qubit", wires=1) res = qml.var(obs) # pylint: disable=use-implicit-booleaness-not-comparison - assert res.shape(dev, Shots(None)) == () - assert res.shape(dev, Shots(100)) == () - - @pytest.mark.parametrize( - "obs", - [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)], - ) - def test_shape_shot_vector(self, obs): - """Test that the shape is correct with the shot vector too.""" - res = qml.var(obs) - shot_vector = (1, 2, 3) - dev = qml.device("default.qubit", wires=3, shots=shot_vector) - assert res.shape(dev, Shots(shot_vector)) == ((), (), ()) + assert res.shape(None, 1) == () + assert res.shape(100, 1) == () @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])]) @pytest.mark.parametrize("shots", [None, 1000, [1000, 1111]]) diff --git a/tests/measurements/test_vn_entropy.py b/tests/measurements/test_vn_entropy.py index 7ab1fbbfc80..bd9b199a12f 100644 --- a/tests/measurements/test_vn_entropy.py +++ b/tests/measurements/test_vn_entropy.py @@ -18,7 +18,7 @@ import pytest import pennylane as qml -from pennylane.measurements import Shots, VnEntropy +from pennylane.measurements import VnEntropy from pennylane.measurements.vn_entropy import VnEntropyMP from pennylane.wires import Wires @@ -109,13 +109,12 @@ def test_properties(self): assert meas.numeric_type == float assert meas.return_type == VnEntropy - @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))]) + @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ())]) def test_shape(self, shots, shape): """Test the ``shape`` method.""" meas = qml.vn_entropy(wires=0) - dev = qml.device("default.qubit", wires=1, shots=shots) - assert meas.shape(dev, Shots(shots)) == shape + assert meas.shape(shots, 1) == shape class TestIntegration: From e84286d2e068b96ac1d897fdb08b692a121d31f2 Mon Sep 17 00:00:00 2001 From: albi3ro Date: Fri, 16 Aug 2024 17:35:56 -0400 Subject: [PATCH 5/5] fix null qubit --- pennylane/devices/null_qubit.py | 89 ++++++++++++++++---------------- tests/devices/test_null_qubit.py | 17 +++--- 2 files changed, 50 insertions(+), 56 deletions(-) diff --git a/pennylane/devices/null_qubit.py b/pennylane/devices/null_qubit.py index 9105edc2cee..7d645b8537f 100644 --- a/pennylane/devices/null_qubit.py +++ b/pennylane/devices/null_qubit.py @@ -22,7 +22,7 @@ from dataclasses import replace from functools import singledispatch from numbers import Number -from typing import Union +from typing import Optional, Union import numpy as np @@ -37,7 +37,6 @@ MeasurementProcess, MeasurementValue, ProbabilityMP, - Shots, StateMP, ) from pennylane.tape import QuantumTape, QuantumTapeBatch @@ -56,57 +55,52 @@ @singledispatch def zero_measurement( - mp: MeasurementProcess, obj_with_wires, shots: Shots, batch_size: int, interface: str + mp: MeasurementProcess, num_device_wires, shots: Optional[int], batch_size: int, interface: str ): """Create all-zero results for various measurement processes.""" - return _zero_measurement(mp, obj_with_wires, shots, batch_size, interface) + return _zero_measurement(mp, num_device_wires, shots, batch_size, interface) -def _zero_measurement(mp, obj_with_wires, shots, batch_size, interface): - shape = mp.shape(obj_with_wires, shots) - if all(isinstance(s, int) for s in shape): - if batch_size is not None: - shape = (batch_size,) + shape - return math.zeros(shape, like=interface, dtype=mp.numeric_type) +def _zero_measurement( + mp: MeasurementProcess, num_device_wires: int, shots: Optional[int], batch_size, interface +): + shape = mp.shape(shots, num_device_wires) if batch_size is not None: - shape = ((batch_size,) + s for s in shape) - return tuple(math.zeros(s, like=interface, dtype=mp.numeric_type) for s in shape) + shape = (batch_size,) + shape + return math.zeros(shape, like=interface, dtype=mp.numeric_type) @zero_measurement.register -def _(mp: ClassicalShadowMP, obj_with_wires, shots, batch_size, interface): - shapes = [mp.shape(obj_with_wires, Shots(s)) for s in shots] +def _(mp: ClassicalShadowMP, num_device_wires, shots: Optional[int], batch_size, interface): if batch_size is not None: # shapes = [(batch_size,) + shape for shape in shapes] raise ValueError( "Parameter broadcasting is not supported with null.qubit and qml.classical_shadow" ) - results = tuple(math.zeros(shape, like=interface, dtype=np.int8) for shape in shapes) - return results if shots.has_partitioned_shots else results[0] + shape = mp.shape(shots, num_device_wires) + return math.zeros(shape, like=interface, dtype=np.int8) @zero_measurement.register -def _(mp: CountsMP, obj_with_wires, shots, batch_size, interface): +def _(mp: CountsMP, num_device_wires, shots, batch_size, interface): outcomes = [] if mp.obs is None and not isinstance(mp.mv, MeasurementValue): - num_wires = len(obj_with_wires.wires) - state = "0" * num_wires - results = tuple({state: math.asarray(s, like=interface)} for s in shots) + state = "0" * num_device_wires + results = {state: math.asarray(shots, like=interface)} if mp.all_outcomes: - outcomes = [f"{x:0{num_wires}b}" for x in range(1, 2**num_wires)] + outcomes = [f"{x:0{num_device_wires}b}" for x in range(1, 2**num_device_wires)] else: outcomes = sorted(mp.eigvals()) # always assign shots to the smallest - results = tuple({outcomes[0]: math.asarray(s, like=interface)} for s in shots) + results = {outcomes[0]: math.asarray(shots, like=interface)} outcomes = outcomes[1:] if mp.all_outcomes else [] if outcomes: zero = math.asarray(0, like=interface) - for res in results: - for val in outcomes: - res[val] = zero + for val in outcomes: + results[val] = zero if batch_size is not None: - results = tuple([r] * batch_size for r in results) - return results[0] if len(results) == 1 else results + results = tuple(results for _ in range(batch_size)) + return results zero_measurement.register(DensityMatrixMP)(_zero_measurement) @@ -114,13 +108,18 @@ def _(mp: CountsMP, obj_with_wires, shots, batch_size, interface): @zero_measurement.register(StateMP) @zero_measurement.register(ProbabilityMP) -def _(mp: Union[StateMP, ProbabilityMP], obj_with_wires, shots, batch_size, interface): - wires = mp.wires or obj_with_wires.wires - state = [1.0] + [0.0] * (2 ** len(wires) - 1) +def _( + mp: Union[StateMP, ProbabilityMP], + num_device_wires: int, + shots: Optional[int], + batch_size, + interface, +): + num_wires = len(mp.wires) or num_device_wires + state = [1.0] + [0.0] * (2**num_wires - 1) if batch_size is not None: state = [state] * batch_size - result = math.asarray(state, like=interface) - return (result,) * shots.num_copies if shots.has_partitioned_shots else result + return math.asarray(state, like=interface) @simulator_tracking @@ -224,26 +223,26 @@ def __init__(self, wires=None, shots=None) -> None: self._debugger = None def _simulate(self, circuit, interface): - shots = circuit.shots - obj_with_wires = self if self.wires else circuit - results = tuple( - zero_measurement(mp, obj_with_wires, shots, circuit.batch_size, interface) - for mp in circuit.measurements - ) - if len(results) == 1: - return results[0] - if shots.has_partitioned_shots: - return tuple(zip(*results)) - return results + num_device_wires = len(self.wires) if self.wires else len(circuit.wires) + results = [] + for s in circuit.shots or [None]: + r = tuple( + zero_measurement(mp, num_device_wires, s, circuit.batch_size, interface) + for mp in circuit.measurements + ) + results.append(r[0] if len(circuit.measurements) == 1 else r) + if circuit.shots.has_partitioned_shots: + return tuple(results) + return results[0] def _derivatives(self, circuit, interface): shots = circuit.shots - obj_with_wires = self if self.wires else circuit + num_device_wires = len(self.wires) if self.wires else len(circuit.wires) n = len(circuit.trainable_params) derivatives = tuple( ( math.zeros_like( - zero_measurement(mp, obj_with_wires, shots, circuit.batch_size, interface) + zero_measurement(mp, num_device_wires, shots, circuit.batch_size, interface) ), ) * n diff --git a/tests/devices/test_null_qubit.py b/tests/devices/test_null_qubit.py index e206dae2948..2af73955322 100644 --- a/tests/devices/test_null_qubit.py +++ b/tests/devices/test_null_qubit.py @@ -466,13 +466,11 @@ def test_counts_wires_batched(self, all_outcomes): dev = NullQubit() result = dev.execute(qs) - print(result) - assert ( - result == tuple([{"00": s, "01": 0, "10": 0, "11": 0}] * 2 for s in qs.shots) - if all_outcomes - else tuple([{"00": s}] * 2 for s in qs.shots) - ) + if all_outcomes: + assert result == tuple(({"00": s, "01": 0, "10": 0, "11": 0},) * 2 for s in qs.shots) + else: + assert result == tuple(({"00": s},) * 2 for s in qs.shots) @pytest.mark.parametrize("all_outcomes", [False, True]) def test_counts_obs(self, all_outcomes): @@ -500,11 +498,10 @@ def test_counts_obs_batched(self, all_outcomes): dev = NullQubit() result = dev.execute(qs) - print(result) assert ( - result == tuple([{-1: s, 1: 0}] * 2 for s in qs.shots) + result == tuple(({-1: s, 1: 0},) * 2 for s in qs.shots) if all_outcomes - else tuple([{-1: s}] * 2 for s in qs.shots) + else tuple(({-1: s},) * 2 for s in qs.shots) ) @@ -929,7 +926,6 @@ def test_vjps_many_tapes_many_results(self): cotangents = [(0.456,), (0.789, 0.123)] actual_grad = dev.compute_vjp([single_meas, multi_meas], cotangents, self.ec) - print(actual_grad) assert actual_grad == ((0.0,), (0.0,)) actual_val, actual_grad = dev.execute_and_compute_vjp( @@ -956,7 +952,6 @@ def test_vjps_integration(self): assert new_ec.use_device_gradient assert new_ec.grad_on_execution - print(actual_grad) assert actual_grad == ((0.0,), (0.0,))