Add dynamic_one_shot tensorflow support + expanded testing

doc/releases/changelog-0.37.0.md

@@ -929,4 +929,4 @@ Kenya Sakka,
 Jay Soni,
 Kazuki Tsuoka,
 Haochen Paul Wang,
-David Wierichs.
+David Wierichs.

doc/releases/changelog-dev.md

@@ -3,7 +3,11 @@
 # Release 0.38.0-dev (development release)
 
 <h3>New features since last release</h3>
-* A new method `process_density_matrix` has been added to the `ProbabilityMP` and `DensityMatrixMP` classes, allowing for more efficient handling of quantum density matrices, particularly with batch processing support. This method simplifies the calculation of probabilities from quantum states represented as density matrices.
+
+* A new method `process_density_matrix` has been added to the `ProbabilityMP` and `DensityMatrixMP`
+  classes, allowing for more efficient handling of quantum density matrices, particularly with batch
+  processing support. This method simplifies the calculation of probabilities from quantum states
+  represented as density matrices.
   [(#5830)](https://github.com/PennyLaneAI/pennylane/pull/5830)
 
 * The `qml.PrepSelPrep` template is added. The template implements a block-encoding of a linear 
@@ -23,6 +27,9 @@
 * `QuantumScript.hash` is now cached, leading to performance improvements.
   [(#5919)](https://github.com/PennyLaneAI/pennylane/pull/5919)
 
+* `qml.dynamic_one_shot` now supports circuits using the `"tensorflow"` interface.
+  [(#5973)](https://github.com/PennyLaneAI/pennylane/pull/5973)
+
 <h3>Breaking changes 💔</h3>
 
 <h3>Deprecations 👋</h3>

pennylane/devices/default_qubit.py

@@ -505,7 +505,6 @@ def preprocess(
             mid_circuit_measurements,
             device=self,
             mcm_config=config.mcm_config,
-            interface=config.interface,
         )
         transform_program.add_transform(
             decompose,

pennylane/devices/execution_config.py

@@ -36,6 +36,9 @@ class MCMConfig:
     be returned. If ``"fill-shots"``, results corresponding to the original number of
     shots will be returned. If not specified, the device will decide which mode to use."""
 
+    interface: Optional[str] = None
+    """Machine learning interface to be used for processing mid-circuit measurement results."""
+
     def __post_init__(self):
         """
         Validate the configured mid-circuit measurement options.

pennylane/devices/preprocess.py

@@ -158,15 +158,15 @@ def mid_circuit_measurements(
     If the tape or device uses finite-shot, use the native implementation (i.e. no transform),
     and use the ``qml.defer_measurements`` transform otherwise.
     """
-
+    # pylint: disable=unused-argument
     if isinstance(mcm_config, dict):
         mcm_config = MCMConfig(**mcm_config)
     mcm_method = mcm_config.mcm_method
     if mcm_method is None:
         mcm_method = "one-shot" if tape.shots else "deferred"
 
     if mcm_method == "one-shot":
-        return qml.dynamic_one_shot(tape, interface=interface)
+        return qml.dynamic_one_shot(tape, interface=mcm_config.interface)
     if mcm_method == "tree-traversal":
         return (tape,), null_postprocessing
     return qml.defer_measurements(tape, device=device)

pennylane/devices/qubit/apply_operation.py

@@ -328,7 +328,7 @@ def apply_mid_measure(
         axis = wire.toarray()[0]
         slices = [slice(None)] * qml.math.ndim(state)
         slices[axis] = 0
-        prob0 = qml.math.norm(state[tuple(slices)]) ** 2
+        prob0 = qml.math.real(qml.math.norm(state[tuple(slices)]) ** 2)
 
         if prng_key is not None:
             # pylint: disable=import-outside-toplevel

pennylane/math/single_dispatch.py

@@ -237,6 +237,7 @@ def _take_autograd(tensor, indices, axis=None):
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "arctan"] = "tensorflow.math"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "arctan2"] = "tensorflow.math"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "mod"] = "tensorflow.math"
+ar.autoray._SUBMODULE_ALIASES["tensorflow", "logical_and"] = "tensorflow.math"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "kron"] = "tensorflow.experimental.numpy"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "moveaxis"] = "tensorflow.experimental.numpy"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "sinc"] = "tensorflow.experimental.numpy"

pennylane/measurements/mid_measure.py

@@ -443,7 +443,7 @@ def __mul__(self, other):
         return self._transform_bin_op(lambda a, b: a * b, other)
 
     def __rmul__(self, other):
-        return self._apply(lambda v: other * v)
+        return self._apply(lambda v: other * qml.math.cast_like(v, other))
 
     def __truediv__(self, other):
         return self._transform_bin_op(lambda a, b: a / b, other)

pennylane/transforms/dynamic_one_shot.py

@@ -124,7 +124,7 @@ def func(x, y):
 
     aux_tapes = [init_auxiliary_tape(t) for t in tapes]
 
-    interface = kwargs.get("interface", None)
+    interface = kwargs.get("interface", None) or "numpy"
 
     def reshape_data(array):
         return qml.math.squeeze(qml.math.vstack(array))
@@ -246,20 +246,27 @@ def measurement_with_no_shots(measurement):
             else np.nan
         )
 
-    interface = interface or qml.math.get_deep_interface(circuit.data)
-    interface = "numpy" if interface == "builtins" else interface
+    results_interface = qml.math.get_deep_interface(results)
+    results_interface = "numpy" if results_interface == "builtins" else results_interface
+    results_interface = "tensorflow" if results_interface == "tf" else results_interface
     active_qjit = qml.compiler.active()
 
     all_mcms = [op for op in aux_tapes[0].operations if is_mcm(op)]
     n_mcms = len(all_mcms)
-    mcm_samples = qml.math.hstack(tuple(res.reshape((-1, 1)) for res in results[-n_mcms:]))
-    mcm_samples = qml.math.array(mcm_samples, like=interface)
+    mcm_samples = qml.math.hstack(
+        tuple(qml.math.reshape(res, (-1, 1)) for res in results[-n_mcms:])
+    )
+    mcm_samples = qml.math.array(mcm_samples, like=results_interface)
     # Can't use boolean dtype array with tf, hence why conditionally setting items to 0 or 1
     has_postselect = qml.math.array(
-        [[int(op.postselect is not None) for op in all_mcms]], like=interface
+        [[int(op.postselect is not None) for op in all_mcms]],
+        like=results_interface,
+        dtype=mcm_samples.dtype,
     )
     postselect = qml.math.array(
-        [[0 if op.postselect is None else op.postselect for op in all_mcms]], like=interface
+        [[0 if op.postselect is None else op.postselect for op in all_mcms]],
+        like=results_interface,
+        dtype=mcm_samples.dtype,
     )
     is_valid = qml.math.all(mcm_samples * has_postselect == postselect, axis=1)
     has_valid = qml.math.any(is_valid)
@@ -276,9 +283,11 @@ def measurement_with_no_shots(measurement):
         if interface != "jax" and m.mv and not has_valid:
             meas = measurement_with_no_shots(m)
         elif m.mv and active_qjit:
-            meas = gather_mcm_qjit(m, mcm_samples, is_valid)  # pragma: no cover
+            meas = gather_mcm_qjit(
+                m, mcm_samples, is_valid, interface=interface
+            )  # pragma: no cover
         elif m.mv:
-            meas = gather_mcm(m, mcm_samples, is_valid)
+            meas = gather_mcm(m, mcm_samples, is_valid, interface=interface)
         elif interface != "jax" and not has_valid:
             meas = measurement_with_no_shots(m)
             m_count += 1
@@ -288,19 +297,22 @@ def measurement_with_no_shots(measurement):
                 # We don't need to cast to arrays when using qml.counts. qml.math.array is not viable
                 # as it assumes all elements of the input are of builtin python types and not belonging
                 # to any particular interface
-                result = qml.math.array(result, like=interface)
+                result = qml.math.array(result, like=results_interface)
             if active_qjit:  # pragma: no cover
                 # `result` contains (bases, counts) need to return (basis, sum(counts)) where `is_valid`
                 # Any row of `result[0]` contains basis, so we return `result[0][0]`
                 # We return the sum of counts (`result[1]`) weighting by `is_valid`, which is `0` for invalid samples
                 if isinstance(m, CountsMP):
                     normalized_meas.append(
-                        (result[0][0], qml.math.sum(result[1] * is_valid.reshape((-1, 1)), axis=0))
+                        (
+                            result[0][0],
+                            qml.math.sum(result[1] * qml.math.reshape(is_valid, (-1, 1)), axis=0),
+                        )
                     )
                     m_count += 1
                     continue
                 result = qml.math.squeeze(result)
-            meas = gather_non_mcm(m, result, is_valid)
+            meas = gather_non_mcm(m, result, is_valid, interface=interface)
             m_count += 1
         if isinstance(m, SampleMP):
             meas = qml.math.squeeze(meas)
@@ -309,7 +321,7 @@ def measurement_with_no_shots(measurement):
     return tuple(normalized_meas) if len(normalized_meas) > 1 else normalized_meas[0]
 
 
-def gather_mcm_qjit(measurement, samples, is_valid):  # pragma: no cover
+def gather_mcm_qjit(measurement, samples, is_valid, interface=None):  # pragma: no cover
     """Process MCM measurements when the Catalyst compiler is active.
 
     Args:
@@ -330,20 +342,20 @@ def gather_mcm_qjit(measurement, samples, is_valid):  # pragma: no cover
         raise LookupError("MCM not found")
     meas = qml.math.squeeze(meas)
     if isinstance(measurement, (CountsMP, ProbabilityMP)):
-        interface = qml.math.get_deep_interface(is_valid)
+        results_interface = qml.math.get_interface(is_valid)
         sum_valid = qml.math.sum(is_valid)
         count_1 = qml.math.sum(meas * is_valid)
         if isinstance(measurement, CountsMP):
-            return qml.math.array([0, 1], like=interface), qml.math.array(
-                [sum_valid - count_1, count_1], like=interface
+            return qml.math.array([0, 1], like=results_interface), qml.math.array(
+                [sum_valid - count_1, count_1], like=results_interface
             )
         if isinstance(measurement, ProbabilityMP):
-            counts = qml.math.array([sum_valid - count_1, count_1], like=interface)
+            counts = qml.math.array([sum_valid - count_1, count_1], like=results_interface)
             return counts / sum_valid
-    return gather_non_mcm(measurement, meas, is_valid)
+    return gather_non_mcm(measurement, meas, is_valid, interface=interface)
 
 
-def gather_non_mcm(measurement, samples, is_valid):
+def gather_non_mcm(measurement, samples, is_valid, interface=None):
     """Combines, gathers and normalizes several measurements with trivial measurement values.
 
     Args:
@@ -364,25 +376,40 @@ def gather_non_mcm(measurement, samples, is_valid):
         if not measurement.all_outcomes:
             tmp = Counter({k: v for k, v in tmp.items() if v > 0})
         return dict(sorted(tmp.items()))
+
+    if (results_interface := qml.math.get_interface(is_valid)) == "tensorflow" and not isinstance(
+        measurement, SampleMP
+    ):
+        # Tensorflow requires arrays that are used for arithmetic with each other to have the
+        # same dtype. We don't cast if measuring samples as float tf.Tensors cannot be used to
+        # index other tf.Tensors (is_valid is used to index valid samples).
+        is_valid = qml.math.cast_like(is_valid, samples)
+
     if isinstance(measurement, ExpectationMP):
         return qml.math.sum(samples * is_valid) / qml.math.sum(is_valid)
     if isinstance(measurement, ProbabilityMP):
-        return qml.math.sum(samples * is_valid.reshape((-1, 1)), axis=0) / qml.math.sum(is_valid)
+        return qml.math.sum(samples * qml.math.reshape(is_valid, (-1, 1)), axis=0) / qml.math.sum(
+            is_valid
+        )
     if isinstance(measurement, SampleMP):
-        is_interface_jax = qml.math.get_deep_interface(is_valid) == "jax"
+        is_interface_jax = interface == "jax"
         if is_interface_jax and samples.ndim == 2:
-            is_valid = is_valid.reshape((-1, 1))
+            is_valid = qml.math.reshape(is_valid, (-1, 1))
         return (
             qml.math.where(is_valid, samples, fill_in_value)
             if is_interface_jax
             else samples[is_valid]
         )
     # VarianceMP
     expval = qml.math.sum(samples * is_valid) / qml.math.sum(is_valid)
+    if results_interface == "tensorflow":
+        # Casting needed for tensorflow
+        samples = qml.math.cast_like(samples, expval)
+        is_valid = qml.math.cast_like(is_valid, expval)
     return qml.math.sum((samples - expval) ** 2 * is_valid) / qml.math.sum(is_valid)
 
 
-def gather_mcm(measurement, samples, is_valid):
+def gather_mcm(measurement, samples, is_valid, interface=None):
     """Combines, gathers and normalizes several measurements with non-trivial measurement values.
 
     Args:
@@ -394,7 +421,7 @@ def gather_mcm(measurement, samples, is_valid):
     Returns:
         TensorLike: The combined measurement outcome
     """
-    interface = qml.math.get_deep_interface(is_valid)
+    results_interface = qml.math.get_deep_interface(is_valid)
     mv = measurement.mv
     # The following block handles measurement value lists, like ``qml.counts(op=[mcm0, mcm1, mcm2])``.
     if isinstance(measurement, (CountsMP, ProbabilityMP, SampleMP)) and isinstance(mv, Sequence):
@@ -403,20 +430,32 @@ def gather_mcm(measurement, samples, is_valid):
         if isinstance(measurement, ProbabilityMP):
             values = [list(m.branches.values()) for m in mv]
             values = list(itertools.product(*values))
-            values = [qml.math.array([v], like=interface) for v in values]
+            values = [
+                qml.math.array([v], like=results_interface, dtype=mcm_samples.dtype) for v in values
+            ]
+            # Need to use boolean functions explicitly as Tensorflow does not allow integer math
+            # on boolean arrays
             counts = [
-                qml.math.sum(qml.math.all(mcm_samples == v, axis=1) * is_valid) for v in values
+                qml.math.count_nonzero(
+                    qml.math.logical_and(qml.math.all(mcm_samples == v, axis=1), is_valid)
+                )
+                for v in values
             ]
-            counts = qml.math.array(counts, like=interface)
+            counts = qml.math.array(counts, like=results_interface)
             return counts / qml.math.sum(counts)
         if isinstance(measurement, CountsMP):
             mcm_samples = [{"".join(str(int(v)) for v in tuple(s)): 1} for s in mcm_samples]
         return gather_non_mcm(measurement, mcm_samples, is_valid)
-    mcm_samples = qml.math.ravel(qml.math.array(mv.concretize(samples), like=interface))
+    mcm_samples = qml.math.ravel(qml.math.array(mv.concretize(samples), like=results_interface))
     if isinstance(measurement, ProbabilityMP):
-        counts = [qml.math.sum((mcm_samples == v) * is_valid) for v in list(mv.branches.values())]
-        counts = qml.math.array(counts, like=interface)
+        # Need to use boolean functions explicitly as Tensorflow does not allow integer math
+        # on boolean arrays
+        counts = [
+            qml.math.count_nonzero(qml.math.logical_and((mcm_samples == v), is_valid))
+            for v in list(mv.branches.values())
+        ]
+        counts = qml.math.array(counts, like=results_interface)
         return counts / qml.math.sum(counts)
     if isinstance(measurement, CountsMP):
         mcm_samples = [{float(s): 1} for s in mcm_samples]
-    return gather_non_mcm(measurement, mcm_samples, is_valid)
+    return gather_non_mcm(measurement, mcm_samples, is_valid, interface=interface)

pennylane/workflow/execution.py

@@ -595,6 +595,9 @@ def cost_fn(params, x):
                 "mcm_method='deferred'."
             )
         config.mcm_config.postselect_mode = "fill-shots"
+    mcm_interface = "tensorflow" if "tf" in mcm_interface else mcm_interface
+    mcm_interface = "jax" if "jax" in mcm_interface else mcm_interface
+    config.mcm_config.interface = mcm_interface
 
     is_gradient_transform = isinstance(gradient_fn, qml.transforms.core.TransformDispatcher)
     transform_program, inner_transform = _make_transform_programs(

pennylane/workflow/qnode.py

@@ -64,7 +64,7 @@ def _get_device_shots(device) -> Shots:
 
 
 def _make_execution_config(
-    circuit: Optional["QNode"], diff_method=None
+    circuit: Optional["QNode"], diff_method=None, mcm_config=None
 ) -> "qml.devices.ExecutionConfig":
     if diff_method is None or isinstance(diff_method, str):
         _gradient_method = diff_method
@@ -76,14 +76,13 @@ def _make_execution_config(
         grad_on_execution = False
     elif grad_on_execution == "best":
         grad_on_execution = None
-    mcm_config = execute_kwargs.get("mcm_config", {})
 
     return qml.devices.ExecutionConfig(
         interface=getattr(circuit, "interface", None),
         gradient_method=_gradient_method,
         grad_on_execution=grad_on_execution,
         use_device_jacobian_product=execute_kwargs.get("device_vjp", False),
-        mcm_config=mcm_config,
+        mcm_config=mcm_config or qml.devices.MCMConfig(),
     )
 
 
@@ -1061,7 +1060,7 @@ def _execution_component(self, args: tuple, kwargs: dict, override_shots) -> qml
 
         if isinstance(self.device, qml.devices.Device):
 
-            config = _make_execution_config(self, self.gradient_fn)
+            config = _make_execution_config(self, self.gradient_fn, mcm_config)
             device_transform_program, config = self.device.preprocess(execution_config=config)
 
             if config.use_device_gradient:
@@ -1075,11 +1074,11 @@ def _execution_component(self, args: tuple, kwargs: dict, override_shots) -> qml
             and self.device.capabilities().get("supports_mid_measure", False)
         )
         if has_mcm_support:
+            mcm_config.interface = qml.math.get_deep_interface(self.tape.data)
             inner_transform_program.add_transform(
                 qml.devices.preprocess.mid_circuit_measurements,
                 device=self.device,
                 mcm_config=mcm_config,
-                interface=self.interface,
             )
             override_shots = 1
         elif hasattr(self.device, "capabilities"):