Skip to content

Commit

Permalink
Merge branch 'master' into eq-updates
Browse files Browse the repository at this point in the history
  • Loading branch information
mudit2812 authored Jul 28, 2023
2 parents e999c4c + 7b03161 commit fa82ce7
Show file tree
Hide file tree
Showing 12 changed files with 755 additions and 128 deletions.
7 changes: 7 additions & 0 deletions doc/releases/changelog-dev.md
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,13 @@
* PennyLane no longer directly relies on `Operator.__eq__`.
[(#4398)](https://github.com/PennyLaneAI/pennylane/pull/4398)

* If no seed is specified on initialization with `DefaultQubit2`, the local random number generator will be
seeded from on the NumPy's global random number generator.
[(#4394)](https://github.com/PennyLaneAI/pennylane/pull/4394)

* The experimental `DefaultQubit2` device now supports computing VJPs and JVPs using the adjoint method.
[(#4374)](https://github.com/PennyLaneAI/pennylane/pull/4374)

<h3>Breaking changes 💔</h3>

* `Operator.expand` now uses the output of `Operator.decomposition` instead of what it queues.
Expand Down
297 changes: 273 additions & 24 deletions pennylane/devices/experimental/default_qubit_2.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
"""

from functools import partial
from numbers import Number
from typing import Union, Callable, Tuple, Optional, Sequence
import concurrent.futures
import os
Expand All @@ -29,9 +30,9 @@

from . import Device
from .execution_config import ExecutionConfig, DefaultExecutionConfig
from ..qubit.simulate import simulate
from ..qubit.simulate import simulate, get_final_state, measure_final_state
from ..qubit.preprocess import preprocess, validate_and_expand_adjoint
from ..qubit.adjoint_jacobian import adjoint_jacobian
from ..qubit.adjoint_jacobian import adjoint_jacobian, adjoint_vjp, adjoint_jvp

Result_or_ResultBatch = Union[Result, ResultBatch]
QuantumTapeBatch = Sequence[QuantumTape]
Expand All @@ -43,12 +44,14 @@
class DefaultQubit2(Device):
"""A PennyLane device written in Python and capable of backpropagation derivatives.
Args:
seed (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
If no value is provided, a default RNG will be used.
Keyword Args:
seed="global" (Union[str, None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng`` or
a request to seed from numpy's global random number generator.
The default, ``seed="global"`` pulls a seed from NumPy's global generator. ``seed=None``
will pull a seed from the OS entropy.
max_workers (int): A ``ProcessPoolExecutor`` executes tapes asynchronously
max_workers=None (int): A ``ProcessPoolExecutor`` executes tapes asynchronously
using a pool of at most ``max_workers`` processes. If ``max_workers`` is ``None``,
only the current process executes tapes. If you experience any
issue, say using JAX, TensorFlow, Torch, try setting ``max_workers`` to ``None``.
Expand Down Expand Up @@ -134,9 +137,10 @@ def name(self):
"""The name of the device."""
return "default.qubit.2"

def __init__(self, seed=None, max_workers=None) -> None:
def __init__(self, seed="global", max_workers=None) -> None:
super().__init__()
self._max_workers = max_workers
seed = np.random.randint(0, high=10000000) if seed == "global" else seed
self._rng = np.random.default_rng(seed)
self._debugger = None

Expand Down Expand Up @@ -169,9 +173,10 @@ def supports_derivatives(
):
return True

if execution_config.gradient_method == "adjoint":
if execution_config.gradient_method == "adjoint" and execution_config.use_device_gradient:
if circuit is None:
return True

return isinstance(validate_and_expand_adjoint(circuit), QuantumScript)

return False
Expand Down Expand Up @@ -264,24 +269,247 @@ def compute_derivatives(
self.tracker.update(derivative_batches=1, derivatives=len(circuits))
self.tracker.record()

if execution_config.gradient_method == "adjoint":
max_workers = self._get_max_workers(execution_config)
if max_workers is None:
res = tuple(adjoint_jacobian(circuit) for circuit in circuits)
else:
vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
exec_map = executor.map(adjoint_jacobian, vanilla_circuits)
res = tuple(circuit for circuit in exec_map)
max_workers = self._get_max_workers(execution_config)
if max_workers is None:
res = tuple(adjoint_jacobian(circuit) for circuit in circuits)
else:
vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
exec_map = executor.map(adjoint_jacobian, vanilla_circuits)
res = tuple(circuit for circuit in exec_map)

# reset _rng to mimic serial behavior
self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
# reset _rng to mimic serial behavior
self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))

return res[0] if is_single_circuit else res
return res[0] if is_single_circuit else res

raise NotImplementedError(
f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
)
def execute_and_compute_derivatives(
self,
circuits: QuantumTape_or_Batch,
execution_config: ExecutionConfig = DefaultExecutionConfig,
):
is_single_circuit = False
if isinstance(circuits, QuantumScript):
is_single_circuit = True
circuits = [circuits]

if self.tracker.active:
for c in circuits:
self.tracker.update(resources=c.specs["resources"])
self.tracker.update(
execute_and_derivative_batches=1,
executions=len(circuits),
derivatives=len(circuits),
)
self.tracker.record()

max_workers = self._get_max_workers(execution_config)
if max_workers is None:
results = tuple(
_adjoint_jac_wrapper(c, rng=self._rng, debugger=self._debugger) for c in circuits
)
results, jacs = tuple(zip(*results))
else:
self._validate_multiprocessing_circuits(circuits)

vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
results = tuple(executor.map(_adjoint_jac_wrapper, vanilla_circuits, seeds))

results, jacs = tuple(zip(*results))

# reset _rng to mimic serial behavior
self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))

return (results[0], jacs[0]) if is_single_circuit else (results, jacs)

def supports_jvp(
self,
execution_config: Optional[ExecutionConfig] = None,
circuit: Optional[QuantumTape] = None,
) -> bool:
"""Whether or not this device defines a custom jacobian vector product.
``DefaultQubit2`` supports backpropagation derivatives with analytic results, as well as
adjoint differentiation.
Args:
execution_config (ExecutionConfig): The configuration of the desired derivative calculation
circuit (QuantumTape): An optional circuit to check derivatives support for.
Returns:
bool: Whether or not a derivative can be calculated provided the given information
"""
return self.supports_derivatives(execution_config, circuit)

def compute_jvp(
self,
circuits: QuantumTape_or_Batch,
tangents: Tuple[Number],
execution_config: ExecutionConfig = DefaultExecutionConfig,
):
is_single_circuit = False
if isinstance(circuits, QuantumScript):
is_single_circuit = True
circuits = [circuits]
tangents = [tangents]

if self.tracker.active:
self.tracker.update(jvp_batches=1, jvps=len(circuits))
self.tracker.record()

max_workers = self._get_max_workers(execution_config)
if max_workers is None:
res = tuple(adjoint_jvp(circuit, tans) for circuit, tans in zip(circuits, tangents))
else:
vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
res = tuple(executor.map(adjoint_jvp, vanilla_circuits, tangents))

# reset _rng to mimic serial behavior
self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))

return res[0] if is_single_circuit else res

def execute_and_compute_jvp(
self,
circuits: QuantumTape_or_Batch,
tangents: Tuple[Number],
execution_config: ExecutionConfig = DefaultExecutionConfig,
):
is_single_circuit = False
if isinstance(circuits, QuantumScript):
is_single_circuit = True
circuits = [circuits]
tangents = [tangents]

if self.tracker.active:
for c in circuits:
self.tracker.update(resources=c.specs["resources"])
self.tracker.update(
execute_and_jvp_batches=1, executions=len(circuits), jvps=len(circuits)
)
self.tracker.record()

max_workers = self._get_max_workers(execution_config)
if max_workers is None:
results = tuple(
_adjoint_jvp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
for c, t in zip(circuits, tangents)
)
results, jvps = tuple(zip(*results))
else:
self._validate_multiprocessing_circuits(circuits)

vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
results = tuple(
executor.map(_adjoint_jvp_wrapper, vanilla_circuits, tangents, seeds)
)

results, jvps = tuple(zip(*results))

# reset _rng to mimic serial behavior
self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))

return (results[0], jvps[0]) if is_single_circuit else (results, jvps)

def supports_vjp(
self,
execution_config: Optional[ExecutionConfig] = None,
circuit: Optional[QuantumTape] = None,
) -> bool:
"""Whether or not this device defines a custom vector jacobian product.
``DefaultQubit2`` supports backpropagation derivatives with analytic results, as well as
adjoint differentiation.
Args:
execution_config (ExecutionConfig): A description of the hyperparameters for the desired computation.
circuit (None, QuantumTape): A specific circuit to check differentation for.
Returns:
bool: Whether or not a derivative can be calculated provided the given information
"""
return self.supports_derivatives(execution_config, circuit)

def compute_vjp(
self,
circuits: QuantumTape_or_Batch,
cotangents: Tuple[Number],
execution_config: ExecutionConfig = DefaultExecutionConfig,
):
is_single_circuit = False
if isinstance(circuits, QuantumScript):
is_single_circuit = True
circuits = [circuits]
cotangents = [cotangents]

if self.tracker.active:
self.tracker.update(vjp_batches=1, vjps=len(circuits))
self.tracker.record()

max_workers = self._get_max_workers(execution_config)
if max_workers is None:
res = tuple(adjoint_vjp(circuit, cots) for circuit, cots in zip(circuits, cotangents))
else:
vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
res = tuple(executor.map(adjoint_vjp, vanilla_circuits, cotangents))

# reset _rng to mimic serial behavior
self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))

return res[0] if is_single_circuit else res

def execute_and_compute_vjp(
self,
circuits: QuantumTape_or_Batch,
cotangents: Tuple[Number],
execution_config: ExecutionConfig = DefaultExecutionConfig,
):
is_single_circuit = False
if isinstance(circuits, QuantumScript):
is_single_circuit = True
circuits = [circuits]
cotangents = [cotangents]

if self.tracker.active:
for c in circuits:
self.tracker.update(resources=c.specs["resources"])
self.tracker.update(
execute_and_vjp_batches=1, executions=len(circuits), vjps=len(circuits)
)
self.tracker.record()

max_workers = self._get_max_workers(execution_config)
if max_workers is None:
results = tuple(
_adjoint_vjp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
for c, t in zip(circuits, cotangents)
)
results, vjps = tuple(zip(*results))
else:
self._validate_multiprocessing_circuits(circuits)

vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
results = tuple(
executor.map(_adjoint_vjp_wrapper, vanilla_circuits, cotangents, seeds)
)

results, vjps = tuple(zip(*results))

# reset _rng to mimic serial behavior
self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))

return (results[0], vjps[0]) if is_single_circuit else (results, vjps)

# pylint: disable=missing-function-docstring
def _get_max_workers(self, execution_config=None):
Expand Down Expand Up @@ -352,3 +580,24 @@ def _validate_multiprocessing_workers(max_workers):
environment variable `{varname}={num_threads_suggest}`.""",
UserWarning,
)


def _adjoint_jac_wrapper(c, rng=None, debugger=None):
state, is_state_batched = get_final_state(c, debugger=debugger)
jac = adjoint_jacobian(c, state=state)
res = measure_final_state(c, state, is_state_batched, rng=rng)
return res, jac


def _adjoint_jvp_wrapper(c, t, rng=None, debugger=None):
state, is_state_batched = get_final_state(c, debugger=debugger)
jvp = adjoint_jvp(c, t, state=state)
res = measure_final_state(c, state, is_state_batched, rng=rng)
return res, jvp


def _adjoint_vjp_wrapper(c, t, rng=None, debugger=None):
state, is_state_batched = get_final_state(c, debugger=debugger)
vjp = adjoint_vjp(c, t, state=state)
res = measure_final_state(c, state, is_state_batched, rng=rng)
return res, vjp
2 changes: 1 addition & 1 deletion pennylane/devices/qubit/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,4 +35,4 @@
from .measure import measure
from .preprocess import preprocess
from .sampling import sample_state, measure_with_samples
from .simulate import simulate
from .simulate import simulate, get_final_state, measure_final_state
Loading

0 comments on commit fa82ce7

Please sign in to comment.