qadence 1.2.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

qadence/backends/horqrux/convert_ops.py

@@ -0,0 +1,273 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from functools import reduce
+from itertools import chain as flatten
+from operator import add
+from typing import Any, Callable, Dict
+
+import jax.numpy as jnp
+from horqrux.gates import NOT, H, I, Rx, Ry, Rz, X, Y, Z
+from horqrux.ops import apply_gate
+from horqrux.types import Gate
+from horqrux.utils import overlap
+from jax import Array
+from jax.tree_util import register_pytree_node_class
+
+from qadence.blocks import (
+    AbstractBlock,
+    AddBlock,
+    ChainBlock,
+    CompositeBlock,
+    KronBlock,
+    ParametricBlock,
+    PrimitiveBlock,
+    ScaleBlock,
+)
+from qadence.operations import CNOT, CRX, CRY, CRZ
+from qadence.types import OpName, ParamDictType
+
+from .config import Configuration
+
+ops_map: Dict[str, Callable] = {
+    OpName.X: X,
+    OpName.Y: Y,
+    OpName.Z: Z,
+    OpName.H: H,
+    OpName.RX: Rx,
+    OpName.RY: Ry,
+    OpName.RZ: Rz,
+    OpName.CRX: Rx,
+    OpName.CRY: Ry,
+    OpName.CRZ: Rz,
+    OpName.CNOT: NOT,
+    OpName.I: I,
+}
+
+supported_gates = list(set(list(ops_map.keys())))
+
+
+@register_pytree_node_class
+@dataclass
+class HorqruxCircuit:
+    operators: list[Gate] = field(default_factory=list)
+
+    def tree_flatten(self) -> tuple[tuple[list[Any]], tuple[()]]:
+        children = (self.operators,)
+        aux_data = ()
+        return (children, aux_data)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data: Any, children: Any) -> Any:
+        return cls(*children, *aux_data)
+
+    def forward(self, state: Array, values: ParamDictType) -> Array:
+        for op in self.operators:
+            state = op.forward(state, values)
+        return state
+
+
+@dataclass
+class HorqruxObservable(HorqruxCircuit):
+    def __init__(self, operators: list[Gate]):
+        super().__init__(operators=operators)
+
+    def _forward(self, state: Array, values: ParamDictType) -> Array:
+        for op in self.operators:
+            state = op.forward(state, values)
+        return state
+
+    def forward(self, state: Array, values: ParamDictType) -> Array:
+        return overlap(state, self._forward(state, values))
+
+
+def convert_observable(
+    block: AbstractBlock, n_qubits: int, config: Configuration
+) -> HorqruxObservable:
+    _ops = convert_block(block, n_qubits, config)
+    return HorqruxObservable(_ops)
+
+
+def convert_block(
+    block: AbstractBlock, n_qubits: int = None, config: Configuration = Configuration()
+) -> list:
+    if n_qubits is None:
+        n_qubits = max(block.qubit_support) + 1
+    ops = []
+    if isinstance(block, CompositeBlock):
+        ops = list(flatten(*(convert_block(b, n_qubits, config) for b in block.blocks)))
+        if isinstance(block, AddBlock):
+            ops = [HorqAddGate(ops)]
+        elif isinstance(block, ChainBlock):
+            ops = [HorqruxCircuit(ops)]
+        elif isinstance(block, KronBlock):
+            if all(
+                [
+                    isinstance(b, ParametricBlock) and not isinstance(b, ScaleBlock)
+                    for b in block.blocks
+                ]
+            ):
+                param_names = [config.get_param_name(b)[0] for b in block.blocks if b.is_parametric]
+                ops = [
+                    HorqKronParametric(
+                        gates=[ops_map[b.name] for b in block.blocks],
+                        target=[b.qubit_support[0] for b in block.blocks],
+                        param_names=param_names,
+                    )
+                ]
+
+            elif all([b.name == "CNOT" for b in block.blocks]):
+                ops = [
+                    HorqKronCNOT(
+                        gates=[ops_map[b.name] for b in block.blocks],
+                        target=[b.qubit_support[1] for b in block.blocks],
+                        control=[b.qubit_support[0] for b in block.blocks],
+                    )
+                ]
+            else:
+                ops = [HorqruxCircuit(ops)]
+
+    elif isinstance(block, CNOT):
+        native_op = ops_map[block.name]
+        ops = [
+            HorqCNOTGate(native_op, block.qubit_support[0], block.qubit_support[1])
+        ]  # in horqrux target and control are swapped
+
+    elif isinstance(block, (CRX, CRY, CRZ)):
+        native_op = ops_map[block.name]
+        param_name = config.get_param_name(block)[0]
+
+        ops = [
+            HorqParametricGate(
+                gate=native_op,
+                qubit=block.qubit_support[1],
+                parameter_name=param_name,
+                control=block.qubit_support[0],
+                name=block.name,
+            )
+        ]
+    elif isinstance(block, ScaleBlock):
+        op = convert_block(block.block, n_qubits, config=config)[0]
+        param_name = config.get_param_name(block)[0]
+        ops = [HorqScaleGate(op, param_name)]
+
+    elif isinstance(block, ParametricBlock):
+        native_op = ops_map[block.name]
+        if len(block.parameters._uuid_dict) > 1:
+            raise NotImplementedError("Only single-parameter operations are supported.")
+        param_name = config.get_param_name(block)[0]
+
+        ops = [
+            HorqParametricGate(
+                gate=native_op,
+                qubit=block.qubit_support[0],
+                parameter_name=param_name,
+            )
+        ]
+
+    elif isinstance(block, PrimitiveBlock):
+        native_op = ops_map[block.name]
+        qubit = block.qubit_support[0]
+        ops = [HorqPrimitiveGate(gate=native_op, qubit=qubit, name=block.name)]
+
+    else:
+        raise NotImplementedError(f"Non-supported operation of type {type(block)}.")
+
+    return ops
+
+
+class HorqPrimitiveGate:
+    def __init__(self, gate: Gate, qubit: int, name: str):
+        self.gates: Gate = gate
+        self.target = qubit
+        self.name = name
+
+    def forward(self, state: Array, values: ParamDictType) -> Array:
+        return apply_gate(state, self.gates(self.target))
+
+    def __repr__(self) -> str:
+        return self.name + f"(target={self.target})"
+
+
+class HorqCNOTGate:
+    def __init__(self, gate: Gate, control: int, target: int):
+        self.gates: Callable = gate
+        self.control: int = control
+        self.target: int = target
+
+    def forward(self, state: Array, values: ParamDictType) -> Array:
+        return apply_gate(state, self.gates(self.target, self.control))
+
+
+class HorqKronParametric:
+    def __init__(self, gates: list[Gate], param_names: list[str], target: list[int]):
+        self.operators: list[Gate] = gates
+        self.target: list[int] = target
+        self.param_names: list[str] = param_names
+
+    def forward(self, state: Array, values: ParamDictType) -> Array:
+        return apply_gate(
+            state,
+            tuple(
+                gate(values[param_name], target)
+                for gate, target, param_name in zip(self.operators, self.target, self.param_names)
+            ),
+        )
+
+
+class HorqKronCNOT(HorqruxCircuit):
+    def __init__(self, gates: list[Gate], target: list[int], control: list[int]):
+        self.operators: list[Gate] = gates
+        self.target: list[int] = target
+        self.control: list[int] = control
+
+    def forward(self, state: Array, values: ParamDictType) -> Array:
+        return apply_gate(
+            state,
+            tuple(
+                gate(target, control)
+                for gate, target, control in zip(self.operators, self.target, self.control)
+            ),
+        )
+
+
+class HorqParametricGate:
+    def __init__(
+        self, gate: Gate, qubit: int, parameter_name: str, control: int = None, name: str = ""
+    ):
+        self.gates: Callable = gate
+        self.target: int = qubit
+        self.parameter: str = parameter_name
+        self.control: int | None = control
+        self.name = name
+
+    def forward(self, state: Array, values: ParamDictType) -> Array:
+        val = jnp.array(values[self.parameter])
+        return apply_gate(state, self.gates(val, self.target, self.control))
+
+    def __repr__(self) -> str:
+        return (
+            self.name
+            + f"(target={self.target}, parameter={self.parameter}, control={self.control})"
+        )
+
+
+class HorqAddGate(HorqruxCircuit):
+    def __init__(self, operations: list[Gate]):
+        self.operators = operations
+        self.name = "Add"
+
+    def forward(self, state: Array, values: ParamDictType = {}) -> Array:
+        return reduce(add, (op.forward(state, values) for op in self.operators))
+
+    def __repr__(self) -> str:
+        return self.name + f"({self.operators})"
+
+
+class HorqScaleGate:
+    def __init__(self, op: Gate, parameter_name: str):
+        self.op = op
+        self.parameter: str = parameter_name
+
+    def forward(self, state: Array, values: ParamDictType) -> Array:
+        return jnp.array(values[self.parameter]) * self.op.forward(state, values)

qadence/backends/jax_utils.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import Any
+
+import jax.numpy as jnp
+from jax import Array, device_get
+from sympy import Expr
+from sympy2jax import SymbolicModule as JaxSympyModule
+from torch import Tensor, cdouble, from_numpy
+
+from qadence.types import ParamDictType
+
+
+def jarr_to_tensor(arr: Array, dtype: Any = cdouble) -> Tensor:
+    return from_numpy(device_get(arr)).to(dtype=dtype)
+
+
+def tensor_to_jnp(tensor: Tensor, dtype: Any = jnp.complex128) -> Array:
+    return (
+        jnp.array(tensor.numpy(), dtype=dtype)
+        if not tensor.requires_grad
+        else jnp.array(tensor.detach().numpy(), dtype=dtype)
+    )
+
+
+def values_to_jax(param_values: dict[str, Tensor]) -> dict[str, Array]:
+    return {key: jnp.array(value.detach().numpy()) for key, value in param_values.items()}
+
+
+def jaxify(expr: Expr) -> JaxSympyModule:
+    return JaxSympyModule(expr)
+
+
+def unhorqify(state: Array) -> Array:
+    """Convert a state of shape [2] * n_qubits + [batch_size] to (batch_size, 2**n_qubits)."""
+    return jnp.ravel(state)
+
+
+def uniform_batchsize(param_values: ParamDictType) -> ParamDictType:
+    max_batch_size = max(p.size for p in param_values.values())
+    batched_values = {
+        k: (v if v.size == max_batch_size else v.repeat(max_batch_size))
+        for k, v in param_values.items()
+    }
+    return batched_values

qadence/backends/pulser/backend.py

@@ -28,7 +28,7 @@ from qadence.noise.protocols import apply_noise
 from qadence.overlap import overlap_exact
 from qadence.register import Register
 from qadence.transpile import transpile
-from qadence.types import BackendName, DeviceType, Endianness
+from qadence.types import BackendName, DeviceType, Endianness, Engine
 
 from .channels import GLOBAL_CHANNEL, LOCAL_CHANNEL
 from .cloud import get_client
@@ -159,6 +159,7 @@ class Backend(BackendInterface):
     with_noise: bool = False
     native_endianness: Endianness = Endianness.BIG
     config: Configuration = field(default_factory=Configuration)
+    engine: Engine = Engine.TORCH
 
     def circuit(self, circ: QuantumCircuit) -> Sequence:
         passes = self.config.transpilation_passes

qadence/backends/pyqtorch/backend.py

@@ -30,7 +30,7 @@ from qadence.transpile import (
     scale_primitive_blocks_only,
     transpile,
 )
-from qadence.types import BackendName, Endianness
+from qadence.types import BackendName, Endianness, Engine
 from qadence.utils import infer_batchsize, int_to_basis
 
 from .config import Configuration, default_passes
@@ -52,6 +52,7 @@ class Backend(BackendInterface):
     with_noise: bool = False
     native_endianness: Endianness = Endianness.BIG
     config: Configuration = field(default_factory=Configuration)
+    engine: Engine = Engine.TORCH
 
     def circuit(self, circuit: QuantumCircuit) -> ConvertedCircuit:
         passes = self.config.transpilation_passes

qadence/backends/pyqtorch/convert_ops.py

@@ -252,12 +252,12 @@ class PyQObservable(Module):
                 convert_block(block, n_qubits, config),
             )
 
-    def forward(self, state: Tensor, values: dict[str, Tensor]) -> Tensor:
-        return pyq.overlap(state, self.operation(state, values))
-
     def run(self, state: Tensor, values: dict[str, Tensor]) -> Tensor:
         return self.operation(state, values)
 
+    def forward(self, state: Tensor, values: dict[str, Tensor]) -> Tensor:
+        return pyq.overlap(state, self.run(state, values))
+
 
 class PyQHamiltonianEvolution(Module):
     def __init__(

qadence/backends/utils.py

@@ -17,8 +17,8 @@ from torch import (
     no_grad,
     rand,
 )
-from torch import flatten as torchflatten
 
+from qadence.types import ParamDictType
 from qadence.utils import Endianness, int_to_basis, is_qadence_shape
 
 FINITE_DIFF_EPS = 1e-06
@@ -92,7 +92,7 @@ def count_bitstrings(sample: Tensor, endianness: Endianness = Endianness.BIG) ->
     )
 
 
-def to_list_of_dicts(param_values: dict[str, Tensor]) -> list[dict[str, float]]:
+def to_list_of_dicts(param_values: ParamDictType) -> list[ParamDictType]:
     if not param_values:
         return [param_values]
 
@@ -119,7 +119,7 @@ def pyqify(state: Tensor, n_qubits: int = None) -> Tensor:
 
 def unpyqify(state: Tensor) -> Tensor:
     """Convert a state of shape [2] * n_qubits + [batch_size] to (batch_size, 2**n_qubits)."""
-    return torchflatten(state, start_dim=0, end_dim=-2).t()
+    return torch.flatten(state, start_dim=0, end_dim=-2).t()
 
 
 def is_pyq_shape(state: Tensor, n_qubits: int) -> bool:
@@ -141,6 +141,11 @@ def validate_state(state: Tensor, n_qubits: int) -> None:
         )
 
 
+def infer_batchsize(param_values: ParamDictType = None) -> int:
+    """Infer the batch_size through the length of the parameter tensors."""
+    return max([len(tensor) for tensor in param_values.values()]) if param_values else 1
+
+
 # The following functions can be used to compute potentially higher order gradients using pyqtorch's
 # native 'jacobian' methods.
 

qadence/blocks/embedding.py

@@ -2,11 +2,10 @@ from __future__ import annotations
 
 from typing import Callable, Iterable, List
 
-import numpy as np
 import sympy
-import sympytorch  # type: ignore [import]
-import torch
-from torch import Tensor
+from numpy import array as nparray
+from numpy import cdouble as npcdouble
+from torch import tensor
 
 from qadence.blocks import (
     AbstractBlock,
@@ -16,9 +15,24 @@ from qadence.blocks.utils import (
     parameters,
     uuid_to_expression,
 )
-from qadence.parameters import evaluate, stringify, torchify
+from qadence.parameters import evaluate, make_differentiable, stringify
+from qadence.types import ArrayLike, DifferentiableExpression, Engine, ParamDictType, TNumber
 
-StrTensorDict = dict[str, Tensor]
+
+def _concretize_parameter(engine: Engine) -> Callable:
+    if engine == Engine.JAX:
+        from jax.numpy import array as jaxarray
+        from jax.numpy import float64 as jaxfloat64
+
+        def concretize_parameter(value: TNumber, trainable: bool = False) -> ArrayLike:
+            return jaxarray([value], dtype=jaxfloat64)
+
+    else:
+
+        def concretize_parameter(value: TNumber, trainable: bool = False) -> ArrayLike:
+            return tensor([value], requires_grad=trainable)
+
+    return concretize_parameter
 
 
 def unique(x: Iterable) -> List:
@@ -26,14 +40,13 @@ def unique(x: Iterable) -> List:
 
 
 def embedding(
-    block: AbstractBlock, to_gate_params: bool = False
-) -> tuple[StrTensorDict, Callable[[StrTensorDict, StrTensorDict], StrTensorDict],]:
-    """Construct embedding function.
+    block: AbstractBlock, to_gate_params: bool = False, engine: Engine = Engine.TORCH
+) -> tuple[ParamDictType, Callable[[ParamDictType, ParamDictType], ParamDictType],]:
+    """Construct embedding function which maps user-facing parameters to either *expression-level*.
 
-    It maps user-facing parameters to either *expression-level*
-    parameters or *gate-level* parameters. The construced embedding function has the signature:
+    parameters or *gate-level* parameters. The constructed embedding function has the signature:
 
-         embedding_fn(params: StrTensorDict, inputs: StrTensorDict) -> StrTensorDict:
+         embedding_fn(params: ParamDictType, inputs: ParamDictType) -> ParamDictType:
 
     which means that it maps the *variational* parameter dict `params` and the *feature* parameter
     dict `inputs` to one new parameter dict `embedded_dict` which holds all parameters that are
@@ -56,6 +69,13 @@ def embedding(
     Returns:
         A tuple with variational parameter dict and the embedding function.
     """
+    concretize_parameter = _concretize_parameter(engine)
+    if engine == Engine.TORCH:
+        cast_dtype = tensor
+    else:
+        from jax.numpy import array
+
+        cast_dtype = array
 
     unique_expressions = unique(expressions(block))
     unique_symbols = [p for p in unique(parameters(block)) if not isinstance(p, sympy.Array)]
@@ -77,16 +97,18 @@ def embedding(
     # we dont need to care about constant symbols if they are contained in an symbolic expression
     # we only care about gate params which are ONLY a constant
 
-    embeddings: dict[sympy.Expr, sympytorch.SymPyModule] = {
-        expr: torchify(expr) for expr in unique_expressions if not expr.is_number
+    embeddings: dict[sympy.Expr, DifferentiableExpression] = {
+        expr: make_differentiable(expr=expr, engine=engine)
+        for expr in unique_expressions
+        if not expr.is_number
     }
 
     uuid_to_expr = uuid_to_expression(block)
 
-    def embedding_fn(params: StrTensorDict, inputs: StrTensorDict) -> StrTensorDict:
-        embedded_params: dict[sympy.Expr, Tensor] = {}
+    def embedding_fn(params: ParamDictType, inputs: ParamDictType) -> ParamDictType:
+        embedded_params: dict[sympy.Expr, ArrayLike] = {}
         for expr, fn in embeddings.items():
-            angle: Tensor
+            angle: ArrayLike
             values = {}
             for symbol in expr.free_symbols:
                 if symbol.name in inputs:
@@ -112,26 +134,26 @@ def embedding(
             embedded_params[e] = params[stringify(e)]
 
         if to_gate_params:
-            gate_lvl_params: StrTensorDict = {}
+            gate_lvl_params: ParamDictType = {}
             for uuid, e in uuid_to_expr.items():
                 gate_lvl_params[uuid] = embedded_params[e]
             return gate_lvl_params
         else:
             return {stringify(k): v for k, v in embedded_params.items()}
 
-    params: StrTensorDict
-    params = {p.name: torch.tensor([p.value], requires_grad=True) for p in trainable_symbols}
+    params: ParamDictType
+    params = {
+        p.name: concretize_parameter(value=p.value, trainable=True) for p in trainable_symbols
+    }
     params.update(
         {
-            stringify(expr): torch.tensor([evaluate(expr)], requires_grad=False)
+            stringify(expr): concretize_parameter(value=evaluate(expr), trainable=False)
             for expr in constant_expressions
         }
     )
     params.update(
         {
-            stringify(expr): torch.tensor(
-                np.array(expr.tolist(), dtype=np.cdouble), requires_grad=False
-            )
+            stringify(expr): cast_dtype(nparray(expr.tolist(), dtype=npcdouble))
             for expr in unique_const_matrices
         }
     )

qadence/blocks/utils.py

@@ -5,7 +5,7 @@ from enum import Enum
 from itertools import chain as _flatten
 from typing import Generator, List, Type, TypeVar, Union, get_args
 
-from sympy import Basic, Expr
+from sympy import Array, Basic, Expr
 from torch import Tensor
 
 from qadence.blocks import (
@@ -503,3 +503,22 @@ def assert_same_block(b1: AbstractBlock, b2: AbstractBlock) -> None:
         ), f"Blocks {b1} and {b2} have differing numbers of parameters."
         for p1, p2 in zip(b1.parameters.expressions(), b2.parameters.expressions()):
             assert p1 == p2
+
+
+def unique_parameters(block: AbstractBlock) -> list[Parameter]:
+    """Return the unique parameters in the block.
+
+    These parameters are the actual user-facing parameters which
+    can be assigned by the user. Multiple gates can contain the
+    same unique parameter
+
+    Returns:
+        list[Parameter]: List of unique parameters in the circuit
+    """
+    symbols = []
+    for p in parameters(block):
+        if isinstance(p, Array):
+            continue
+        elif not p.is_number and p not in symbols:
+            symbols.append(p)
+    return symbols

qadence/circuit.py

@@ -5,10 +5,10 @@ from itertools import chain as flatten
 from pathlib import Path
 from typing import Iterable
 
-from sympy import Array, Basic
+from sympy import Basic
 
 from qadence.blocks import AbstractBlock, AnalogBlock, CompositeBlock, chain
-from qadence.blocks.utils import parameters, primitive_blocks
+from qadence.blocks.utils import parameters, primitive_blocks, unique_parameters
 from qadence.parameters import Parameter
 from qadence.register import Register
 
@@ -88,13 +88,7 @@ class QuantumCircuit:
         Returns:
             list[Parameter]: List of unique parameters in the circuit
         """
-        symbols = []
-        for p in parameters(self.block):
-            if isinstance(p, Array):
-                continue
-            elif not p.is_number and p not in symbols:
-                symbols.append(p)
-        return symbols
+        return unique_parameters(self.block)
 
     @property
     def num_unique_parameters(self) -> int: