Merge pull request #13 from jaeyoo/pavan/control_pqc_gpu

`ControlledPQC` cuquantum support
jaeyoo · May 1, 2023 · 1b10b08 · 1b10b08
2 parents 13ac9dc + 0683d1c
commit 1b10b08
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/tensorflow_quantum/python/layers/high_level/controlled_pqc.py b/tensorflow_quantum/python/layers/high_level/controlled_pqc.py
@@ -128,6 +128,7 @@ def __init__(self,
                  *,
                  repetitions=None,
                  backend='noiseless',
+                 use_cuquantum=False,
                  differentiator=None,
                  **kwargs):
         """Instantiate this layer.
@@ -153,6 +154,8 @@ def __init__(self,
             `sampled_based` is True or it must inherit
             `cirq.sim.simulator.SimulatesExpectationValues` if `sample_based` is
             False.
+        use_cuquantum: Optional Python `bool` indicating whether or not to use 
+            GPU ops
         differentiator: Optional `tfq.differentiator` object to specify how
             gradients of `model_circuit` should be calculated.
         """
@@ -235,10 +238,12 @@ def __init__(self,
 
         if self._analytic:
             self._layer = expectation.Expectation(backend=backend,
-                                                  differentiator=differentiator)
+                                                  differentiator=differentiator,
+                                                  use_cuquantum=use_cuquantum)
         else:
             self._layer = sampled_expectation.SampledExpectation(
-                backend=backend, differentiator=differentiator)
+                backend=backend, differentiator=differentiator,
+                use_cuquantum=use_cuquantum)
 
         self._append_layer = elementary.AddCircuit()