from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import cached_property
from inspect import signature
from typing import Callable, List, Optional, Tuple, Union
import numpy as np
from qibo import Circuit
from qiboml import ndarray
from qiboml.models.encoding import QuantumEncoding
[docs]@dataclass
class CircuitTracer(ABC):
"""Wrapper to the circuit structure that takes care of tracing the circuit construction.
In particular, it computes both the Jacobian wrt the parameters and the inputs of the operations
that construct the circuits.
"""
circuit_structure: List[Union[Circuit, QuantumEncoding, Callable]]
derivation_mode: str = "forward"
@property
@abstractmethod
def engine(self): # pragma: no cover
"""The calculation engine used."""
pass
[docs] @staticmethod
@abstractmethod
def jacfwd(f: Callable, argnums: Union[int, Tuple[int]]): # pragma: no cover
"""The jacobian functional in forward derivation mode."""
pass
[docs] @staticmethod
@abstractmethod
def jacrev(f: Callable, argnums: Union[int, Tuple[int]]): # pragma: no cover
"""The jacobian functional in reverse derivation mode."""
pass
[docs] def _compute_jacobian_functional(
self, circuit: Union[QuantumEncoding, Callable]
) -> Callable:
"""Compute the jacobian functional for the input function.
Args:
circuit (Callable): the input functions that build the circuit.
Returns:
(Callable) the jacobian functional
"""
_tmp_circuit = circuit
def build(x):
is_encoding = isinstance(_tmp_circuit, QuantumEncoding)
if is_encoding:
circuit = _tmp_circuit(x)
else:
# this is needed for symbolic execution with tf
x = [x[i] for i in range(len(x))]
circuit = _tmp_circuit(*x)
out = self.engine.stack(
[
p
for pars in circuit.get_parameters(
include_not_trainable=is_encoding
)
for p in pars
]
)
return self.engine.reshape(out, (-1, 1))
jac = self.jacfwd if self.derivation_mode == "forward" else self.jacrev
return jac(build, argnums=0)
@cached_property
def jacobian_functionals(
self,
) -> dict[int, Callable]:
"""The dictionary containing the jacobian functionals of each element composing
circuit structure.
"""
jacobians = {}
for circ in self.circuit_structure:
if isinstance(circ, Circuit):
jacobians[id(circ)] = self.identity
else:
jacobians[id(circ)] = self._compute_jacobian_functional(circ)
return jacobians
[docs] def nonzero(self, array: ndarray) -> ndarray:
"""The numpy-like np.nonzero function of the current engine."""
return self.engine.nonzero(array)
[docs] def _build_parameters_map(self, jacobian: ndarray) -> dict[int, Tuple[int]]:
"""Construct the mapping between independent and dependent parameters.
In detail, the index of each independent parameter is mapped to the indices
of all the dependent parameters that originated from it. This is particularly
useful for the jacobian of encoders, as it provides a map from the inputs to
the actual gates where they are encoded in the circuit.
"""
par_map = {}
for i, row in enumerate(jacobian):
for j in self.nonzero(row):
j = int(j)
if j in par_map:
par_map[j] += (i,)
else:
par_map[j] = (i,)
return par_map
[docs] def trace(
self, f: Union[Callable, QuantumEncoding], params: ndarray
) -> Tuple[ndarray, dict[int, Tuple[int]]]:
"""Trace the construction of a circuit through the function `f` with given
parameters `params`.
Args:
f (Callable): the function that builds the circuit, either a custom user-defined
function or an encoder.
params (ndarray): the parameters assigned to the built circuit.
Returns:
(Tuple(ndarray, dict)) the computed jacobian and the mapping between the independent
and dependent parameters of `f`.
"""
# we always assume the input is a 1-dim array, even for encodings
# thus the jacobian is always a matrix
jac = self.engine.reshape(
self.jacobian_functionals[id(f)](params), (-1, params.shape[0])
)
par_map = self._build_parameters_map(jac)
return jac, par_map
@cached_property
def is_encoding_differentiable(self) -> bool:
"""Check if all the encoders in the circuit structure are differentiable."""
diff_encodings = [
circ.differentiable
for circ in self.circuit_structure
if isinstance(circ, QuantumEncoding)
]
if len(diff_encodings) == 0:
return False
# all the encodings must be differentiable
# open to debate, not the only possible choice
return all(diff_encodings)
[docs] @abstractmethod
def requires_gradient(self, x: ndarray) -> bool: # pragma: no cover
"""Check whether the input array needs gradients to be calculated for it."""
pass
[docs] def _build_from_encoding(
self, encoding: QuantumEncoding, x: ndarray, trace: bool = True
) -> Circuit:
"""Build the circuit starting from an encoder.
Args:
encoding (QuantumEncoding): the encoder.
x (ndarray): the input data.
trace (bool): whether to trace the construction and thus also
calculate the jacobian. Defaults to ``True``.
Returns:
(Circuit | Tuple(ndarray, dict, Circuit)) the built circuit or the
tuple: jacobian wrt inputs, input to gate map and circuit.
"""
circuit = encoding(x)
if trace:
if self.is_encoding_differentiable and self.requires_gradient(x):
if len(x.shape) > 1:
x = x[0]
return *self.trace(encoding, x), circuit
return None, None, circuit
return circuit
[docs] def _build_from_circuit(
self, circuit: Circuit, params: ndarray, trace: bool = True
) -> Circuit:
"""Build the circuit starting from a circuit. In practice the given
parameters `params` are set into the circuit only.
Args:
circuit (Circuit): the circuit.
params (ndarray): the parameters to set.
trace (bool): whether to trace the construction and thus also
calculate the jacobian. Defaults to ``True``.
Returns:
(Circuit | Tuple(ndarray, dict, Circuit)) the built circuit or the
tuple: jacobian wrt parameters, input to gate map and circuit.
"""
circuit.set_parameters(params)
if trace:
jacobian = self.identity(
params.shape[0],
dtype=self._get_dtype(params),
device=self._get_device(params),
)
# all the circuit parameters are considered independent
par_map = {i: (i,) for i in range(params.shape[0])}
return jacobian, par_map, circuit
return circuit
[docs] def _build_from_callable(
self, f: Callable, params: ndarray, trace: bool = True
) -> Circuit:
"""Build the circuit starting from a custom callable.
Args:
f (Callable): the custom callable which builds the circuit.
params (ndarray): the parameters used for construction.
trace (bool): whether to trace the construction and thus also
calculate the jacobian. Defaults to ``True``.
Returns:
(Circuit | Tuple(ndarray, dict, Circuit)) the built circuit or the
tuple: jacobian wrt parameters, input to gate map and circuit.
"""
circuit = f(*params)
if trace:
return *self.trace(f, params), circuit
return circuit
[docs] @staticmethod
def _get_device(array: ndarray):
"""Extract the device of the input array."""
return array.device
[docs] @staticmethod
def _get_dtype(array: ndarray):
"""Extract the dtype of the input array."""
return array.dtype
[docs] @abstractmethod
def identity(self, dim: int, dtype, device) -> ndarray: # pragma: no cover
"""The numpy-like np.eye function of the current engine."""
pass
[docs] @abstractmethod
def zeros(
self, shape: Union[int, Tuple[int]], dtype, device
) -> ndarray: # pragma: no cover
"""The numpy-like np.zeros function of the current engine."""
pass
[docs] def fill_jacobian(
self,
jacobian: ndarray,
row_span: Tuple[int, int],
col_span: Tuple[int, int],
values: ndarray,
) -> ndarray:
"""Fill the input jacobian in the span defined by (row_span, col_span) with
the given values. This is mostly here to be overwritten by engines that do not
allow for direct item assignment.
Args:
jacobian (ndarray): the jacobian to fill.
row_span (Tuple[int, int]): the span of row indices.
col_span (Tuple[int, int]): the span of column indices.
values (ndarray): the values to insert.
Returns:
(ndarray) the filled jacobian.
"""
jacobian[row_span[0] : row_span[1], col_span[0] : col_span[1]] = values
return jacobian
def __call__(
self, params: ndarray, x: Optional[ndarray] = None
) -> Tuple[Circuit, Optional[ndarray], ndarray, Optional[dict]]:
"""Construct the circuit defined by the internal ``circuit_structure`` and compute
the jacobian wrt to the parameters `params` and inputs `x` of the construction process.
Args:
params (ndarray): the parameters to construct the circuit with.
x (ndarray, optional): the inputs to construct the circuit with.
Returns:
(Tuple(Circuit, ndarray, ndarray, dict)) the constructed circuit, jacobian wrt inputs,
jacobian wrt parameters and mapping of inputs ``x`` to the corresponding gate in the built circuit.
"""
if (
any(isinstance(circ, QuantumEncoding) for circ in self.circuit_structure)
and x is None
):
raise ValueError(
"x cannot be None when encoding layers are present in the circuit structure."
)
# the complete circuit
circuit = None
# the jacobian for each sub-circuit
jacobians = []
jacobians_wrt_inputs = []
input_to_gate_map = {}
index = 0
for circ in self.circuit_structure:
if isinstance(circ, QuantumEncoding):
# encoders do not have parameters
nparams = 0
jacobian, input_map, circ = self._build_from_encoding(
circ, x, trace=True
)
if input_map is not None:
# update the input_map to the index of the global circuit
input_map = {
inp: tuple(i + index for i in indices)
for inp, indices in input_map.items()
}
# update the global map
for inp, indices in input_map.items():
if inp in input_to_gate_map:
input_to_gate_map[inp] += indices
else:
input_to_gate_map[inp] = indices
jacobians_wrt_inputs.append(jacobian)
elif isinstance(circ, Circuit):
nparams = len(circ.get_parameters())
jacobian, _, circ = self._build_from_circuit(
circ, params[index : index + nparams], trace=True
)
jacobians.append(jacobian)
else:
param_dict = signature(circ).parameters
nparams = len(param_dict)
jacobian, _, circ = self._build_from_callable(
circ, params[index : index + nparams], trace=True
)
jacobians.append(jacobian)
index += nparams
if circuit is None:
circuit = circ
else:
circuit += circ
total_dim = tuple(sum(np.array(j.shape) for j in jacobians))
# build the global jacobian
J = self.zeros(
total_dim, dtype=self._get_dtype(params), device=self._get_device(params)
)
position = np.array([0, 0])
# insert each sub-jacobian in the total one
for j in jacobians:
shape = np.array(j.shape)
interval = tuple(zip(position, shape + position))
J = self.fill_jacobian(J, interval[0], interval[1], j)
# direct assignment works with torch/numpy only
# J[interval[0][0] : interval[0][1], interval[1][0] : interval[1][1]] = j
position += shape
jacobians_wrt_inputs = (
None
if len(jacobians_wrt_inputs) == 0 or jacobians_wrt_inputs[0] is None
else self.engine.vstack(jacobians_wrt_inputs)
)
if len(input_to_gate_map) == 0:
input_to_gate_map = None
return circuit, jacobians_wrt_inputs, J, input_to_gate_map
[docs] def build_circuit(self, params: ndarray, x: Optional[ndarray] = None) -> Circuit:
"""Construct the circuit defined by the internal ``circuit_structure`` without
worrying about tracing the operations.
Args:
params (ndarray): the parameters to construct the circuit with.
x (ndarray, optional): the inputs to construct the circuit with.
Returns:
(Circuit) the constructed circuit.
"""
circuit = None
index = 0
for circ in self.circuit_structure:
if isinstance(circ, QuantumEncoding):
nparams = 0
circ = self._build_from_encoding(circ, x, trace=False)
elif isinstance(circ, Circuit):
nparams = len(circ.get_parameters())
circ = self._build_from_circuit(
circ, params[index : index + nparams], trace=False
)
else:
param_dict = signature(circ).parameters
nparams = len(param_dict)
circ = self._build_from_callable(
circ, params[index : index + nparams], trace=False
)
index += nparams
if circuit is None:
circuit = circ
else:
circuit += circ
return circuit