[Return-types Fix] Update mitigate_with_zne transform (#3415)

* Update mitigate for new returns

* Add test for coverage (when mitiq not installed)

* CMP #8: Add `Expectation` MP class (#3343)

* Add process abstract method to MeasurementProcess class

* Add TODO comment.

* Add docstring

* Add _Probability class

* Add _Sample class

* Add process method to ClassicalShadow class

* Add tests

* Remove r_dtype from argument

* Fix coverage

* TODO

* TODO

* Fix test

* Fix equal

* Fix test

* Fix equal

* Support python3.7

* Revert

* Skip python@3.7 test

* Improve sample code

* Coverage

* Add process_state method

* Add process_state method

* Add tests

* Improve tests

* Improve tests

* Add static method

* Improve code

* Add process_state method

* Remove abstract

* Coverage

* Improve code

* Add StateMeasurement and SampleMeasurement classes

* Change inheritance

* Change inheritance

* Revert

* Add changelog entry

* Add changelog entry

* Add changelog entry

* Small fix

* Small fix

* Remove dead code

* Coverage

* Clean up tests

* Fix test

* Coverage

* Add changelog entry

* Add _Expectation class

* Add changelog entry

* Coverage

* Address comments

* Change state type

* Remove error

* Small fix

* Fix types

* Remove numpy usage

* Fix types

* Fix wire_map

* Fix wire_map

* Add wire_order to process_samples

* Docstring

* Small fix

* Add wire_order to process_samples

* Add wire_order to process_samples

* Change tests

* Add wire_order to process_samples

* Add breaking change message

* Change tests

* Add wire_order to process_samples

* Change tests

* Fix tests

* Fix tests

* Update tests/measurements/test_expval.py

* Use tol

* Improve code

* Update doc/releases/changelog-dev.md

* Group lines

* CMP #9: Add `Var` MP class (#3312)

* Add process abstract method to MeasurementProcess class

* Add TODO comment.

* Add docstring

* Add _Probability class

* Add _Sample class

* Add process method to ClassicalShadow class

* Add tests

* Remove r_dtype from argument

* Fix coverage

* TODO

* TODO

* Fix test

* Fix equal

* Fix test

* Fix equal

* Support python3.7

* Revert

* Skip python@3.7 test

* Improve sample code

* Coverage

* Add process_state method

* Add process_state method

* Add tests

* Improve tests

* Improve tests

* Add static method

* Improve code

* Add process_state method

* Remove abstract

* Coverage

* Add _Var MP Class

* Improve code

* Fix tests

* Add StateMeasurement and SampleMeasurement classes

* Change inheritance

* Change inheritance

* Change inheritance

* Revert

* Add changelog entry

* Add changelog entry

* Add changelog entry

* Small fix

* Small fix

* Small fix

* Remove dead code

* Coverage

* Clean up tests

* Small fix

* Coverage

* Coverage

* Fix test

* Coverage

* Add changelog entry

* Address comments

* Change state type

* Remove error

* Small fix

* Fix types

* Remove numpy usage

* Fix types

* Fix wire_map

* Fix wire_map

* Add wire_order to process_samples

* Docstring

* Small fix

* Add wire_order to process_samples

* Add wire_order to process_samples

* Change tests

* Add wire_order to process_samples

* Add breaking change message

* Change tests

* Add wire_order to process_samples

* Change tests

* Update doc/releases/changelog-dev.md

* refactor (_Variance): ♻️ Improve code

* Remove pylint

* test (var): 🧪 Use tol.

* refactor (CMP): ♻️ Move  into  class

Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca>

* Deprecate transforms.measurement_grouping (#3417)

* Add UserWarning regarding deprecation

* Add test for warning

* Update deprecation docs

* add missing blank line

* update changelog

* update changelog

* Add module documentation and explanations for the `queuing` module (#3268)

* add queuing module explanation

* skip documentation of imported objects

* Apply suggestions from code review

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* respond to feedback

* Apply suggestions from code review

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* thought I already made these changes.  remove copy and queue discussion

* more info

* changelog

* Update pennylane/queuing.py

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* Use `convert_like` in `hamiltonian_expand` (#3339)

* covert_like

* assert the types

* minor docs

* Update output type check

* No need to check if len(res) > 0: we should definitely have results if the Hamiltonian has terms in it

* Validate `hamiltonian_expand`: don't allow Hamiltonian with no terms (#3411)

* validate input Hamiltonian

* no thewalrus just yet

* Revert "no thewalrus just yet"

This reverts commit baae20c.

* fix: TF ResourceVariable has no len attr, can't take the len of coeffs and ops; check explicit value comparison

* Revert "fix: TF ResourceVariable has no len attr, can't take the len of coeffs and ops; check explicit value comparison"

This reverts commit f5e5784.

* use qml.math.shape

* correct shape

* trigger CI

* trigger CI

* changelog

Co-authored-by: Albert Mitjans <a.mitjanscoma@gmail.com>

* Add control Hadamard gate: CH (#3408)

* Add CH

* Fix test QFT

* Add changelog

* Typo

* Update tests/ops/qubit/test_non_parametric_ops.py

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* change test

* Add hermitian test

* Add attribute and commuting

* Missing test

* typo

* Update pennylane/ops/qubit/non_parametric_ops.py

Co-authored-by: Christina Lee <christina@xanadu.ai>

* Fix label test

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
Co-authored-by: Christina Lee <christina@xanadu.ai>

* qml.equal for Hamiltonians and Tensors (#3390)

* use singledispatch

* Different types don't match

* fix bug

* rearrange

* clean up

* add NotImplementedError test

* update docstrings

* docstrings

* tests

* fix bugs from merge conflict

* update tests

* docs and changelog

* Update pennylane/ops/functions/equal.py

Co-authored-by: Christina Lee <christina@xanadu.ai>

* Doc updates

Co-authored-by: antalszava <antalszava@gmail.com>

* Remove extra indent

Co-authored-by: antalszava <antalszava@gmail.com>

* Add specific handling for Hermitians

* Split comparing Hamiltonian and Tensor into seperate functions

* docstring for _obs_comparison_data

* add check for op2 being Observable

* Add function for PauliX/Y/Z comparisons

* update tests

* Default to original qml.equal logic if no special case

* Update changelog

* Fix tests

* generalize _equal_pauli to other Operation/Observables

* rearrange to simplify

* update tests

* add tests

* combine _equal_tensor and _obs_comparison_data

* add NotImplementedError

* update tests

Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: antalszava <antalszava@gmail.com>

* Add control CZ gate: CCZ (#3410)

* Add CCZ

* add decomp test

* Add changelog

* Typo

* update control test

* Trigger CI

* Fix test

* Update tests/gate_data.py

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* black apply

* Black

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

* Add changelog entry

* Don't stack results

* Fix and add tests

Co-authored-by: Albert Mitjans <a.mitjanscoma@gmail.com>
Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca>
Co-authored-by: lillian542 <38584660+lillian542@users.noreply.github.com>
Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
Co-authored-by: antalszava <antalszava@gmail.com>
Co-authored-by: Romain Moyard <rmoyard@gmail.com>

doc/releases/changelog-dev.md

@@ -300,6 +300,10 @@
 * Updated `qml.transforms.split_non_commuting` to support the new return types.
   [#3414](https://github.com/PennyLaneAI/pennylane/pull/3414)
 
+* Updated `qml.transforms.mitigate_with_zne` to support the new return types.
+  [#3415](https://github.com/PennyLaneAI/pennylane/pull/3415)
+
+
 <h3>Breaking changes</h3>
 
 * The `log_base` attribute has been moved from `MeasurementProcess` to the new `_VnEntropy` and
@@ -363,6 +367,7 @@ Deprecations cycles are tracked at [doc/developement/deprecations.rst](https://d
   wanted seed.
   [(#3388)](https://github.com/PennyLaneAI/pennylane/pull/3388)
 
+
 <h3>Documentation</h3>
 
 * Adds developer documentation for the queuing module.

pennylane/transforms/mitigate.py

@@ -528,6 +528,12 @@ def circuit(w1, w2):
 
     def processing_fn(results):
         """Maps from input tape executions to an error-mitigated estimate"""
+        if qml.active_return():
+            for i, tape in enumerate(out_tapes):
+                # stack the results if there are multiple measurements
+                # this will not create ragged arrays since only expval measurements are allowed
+                if len(tape.observables) > 1:
+                    results[i] = qml.math.stack(results[i])
 
         # Averaging over reps_per_factor repetitions
         results_flattened = []
@@ -538,6 +544,14 @@ def processing_fn(results):
 
         extrapolated = extrapolate(scale_factors, results_flattened, **extrapolate_kwargs)
 
+        if qml.active_return():
+            extrapolated = extrapolated[0] if shape(extrapolated) == (1,) else extrapolated
+
+            # unstack the results in the case of multiple measurements
+            return (
+                extrapolated if shape(extrapolated) == () else tuple(qml.math.unstack(extrapolated))
+            )
+
         return extrapolated[0] if shape(extrapolated) == (1,) else extrapolated
 
     return out_tapes, processing_fn