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

qadence/engines/differentiable_backend.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections import Counter
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from qadence.backend import Backend, Converted, ConvertedCircuit, ConvertedObservable
+from qadence.blocks import AbstractBlock, PrimitiveBlock
+from qadence.blocks.utils import uuid_to_block
+from qadence.circuit import QuantumCircuit
+from qadence.measurements import Measurements
+from qadence.mitigations import Mitigations
+from qadence.noise import Noise
+from qadence.types import ArrayLike, DiffMode, Endianness, Engine, ParamDictType
+
+
+@dataclass(frozen=True, eq=True)
+class DifferentiableBackend(ABC):
+    """The abstract class which wraps any (non)-natively differentiable QuantumBackend.
+
+       in an automatic differentiation engine.
+
+    Arguments:
+        backend: An instance of the QuantumBackend type perform execution.
+        engine: Which automatic differentiation engine the QuantumBackend runs on.
+        diff_mode: A differentiable mode supported by the differentiation engine.
+    """
+
+    backend: Backend
+    engine: Engine
+    diff_mode: DiffMode
+
+    # TODO: Add differentiable overlap calculation
+    _overlap: Callable = None  # type: ignore [assignment]
+
+    def sample(
+        self,
+        circuit: ConvertedCircuit,
+        param_values: ParamDictType = {},
+        n_shots: int = 100,
+        state: ArrayLike | None = None,
+        noise: Noise | None = None,
+        mitigation: Mitigations | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> list[Counter]:
+        """Sample bitstring from the registered circuit.
+
+        Arguments:
+            circuit: A backend native quantum circuit to be executed.
+            param_values: The values of the parameters after embedding
+            n_shots: The number of shots. Defaults to 1.
+            state: Initial state.
+            noise: A noise model to use.
+            mitigation: A mitigation protocol to apply to noisy samples.
+            endianness: Endianness of the resulting bitstrings.
+
+        Returns:
+            An iterable with all the sampled bitstrings
+        """
+
+        return self.backend.sample(
+            circuit=circuit,
+            param_values=param_values,
+            n_shots=n_shots,
+            state=state,
+            noise=noise,
+            mitigation=mitigation,
+            endianness=endianness,
+        )
+
+    def run(
+        self,
+        circuit: ConvertedCircuit,
+        param_values: ParamDictType = {},
+        state: ArrayLike | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> ArrayLike:
+        """Run on the underlying backend."""
+        return self.backend.run(
+            circuit=circuit, param_values=param_values, state=state, endianness=endianness
+        )
+
+    @abstractmethod
+    def expectation(
+        self,
+        circuit: ConvertedCircuit,
+        observable: list[ConvertedObservable] | ConvertedObservable,
+        param_values: ParamDictType = {},
+        state: ArrayLike | None = None,
+        measurement: Measurements | None = None,
+        noise: Noise | None = None,
+        mitigation: Mitigations | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> Any:
+        """Compute the expectation value of the `circuit` with the given `observable`.
+
+        Arguments:
+            circuit: A converted circuit as returned by `backend.circuit`.
+            observable: A converted observable as returned by `backend.observable`.
+            param_values: _**Already embedded**_ parameters of the circuit. See
+                [`embedding`][qadence.blocks.embedding.embedding] for more info.
+            state: Initial state.
+            measurement: Optional measurement protocol. If None, use
+                exact expectation value with a statevector simulator.
+            noise: A noise model to use.
+            mitigation: The error mitigation to use.
+            endianness: Endianness of the resulting bit strings.
+        """
+        raise NotImplementedError(
+            "A DifferentiableBackend needs to override the expectation method."
+        )
+
+    def default_configuration(self) -> Any:
+        return self.backend.default_configuration()
+
+    def circuit(self, circuit: QuantumCircuit) -> ConvertedCircuit:
+        if self.diff_mode == DiffMode.GPSR:
+            parametrized_blocks = list(uuid_to_block(circuit.block).values())
+            non_prim_blocks = filter(
+                lambda b: not isinstance(b, PrimitiveBlock), parametrized_blocks
+            )
+            if len(list(non_prim_blocks)) > 0:
+                raise ValueError(
+                    "The circuit contains non-primitive blocks that are currently\
+                    not supported by the PSR differentiable mode."
+                )
+        return self.backend.circuit(circuit)
+
+    def observable(self, observable: AbstractBlock, n_qubits: int) -> ConvertedObservable:
+        if self.diff_mode != DiffMode.AD and observable is not None:
+            msg = (
+                f"Differentiation mode '{self.diff_mode}' does not support parametric observables."
+            )
+            if isinstance(observable, list):
+                for obs in observable:
+                    if obs.is_parametric:
+                        raise ValueError(msg)
+            else:
+                if observable.is_parametric:
+                    raise ValueError(msg)
+        return self.backend.observable(observable, n_qubits)
+
+    def convert(
+        self,
+        circuit: QuantumCircuit,
+        observable: list[AbstractBlock] | AbstractBlock | None = None,
+    ) -> Converted:
+        return self.backend.convert(circuit, observable)
+
+    def assign_parameters(self, circuit: ConvertedCircuit, param_values: ParamDictType) -> Any:
+        return self.backend.assign_parameters(circuit, param_values)

qadence/engines/jax/__init__.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from jax import config
+
+from .differentiable_backend import DifferentiableBackend
+from .differentiable_expectation import DifferentiableExpectation
+
+config.update("jax_enable_x64", True)

qadence/engines/jax/differentiable_backend.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from qadence.backend import Backend, ConvertedCircuit, ConvertedObservable
+from qadence.engines.differentiable_backend import (
+    DifferentiableBackend as DifferentiableBackendInterface,
+)
+from qadence.engines.jax.differentiable_expectation import DifferentiableExpectation
+from qadence.measurements import Measurements
+from qadence.mitigations import Mitigations
+from qadence.noise import Noise
+from qadence.types import ArrayLike, DiffMode, Endianness, Engine, ParamDictType
+
+
+class DifferentiableBackend(DifferentiableBackendInterface):
+    """A class which wraps a QuantumBackend with the automatic differentation engine JAX.
+
+    Arguments:
+        backend: An instance of the QuantumBackend type perform execution.
+        diff_mode: A differentiable mode supported by the differentiation engine.
+        **psr_args: Arguments that will be passed on to `DifferentiableExpectation`.
+    """
+
+    def __init__(
+        self,
+        backend: Backend,
+        diff_mode: DiffMode = DiffMode.AD,
+        **psr_args: int | float | None,
+    ) -> None:
+        super().__init__(backend=backend, engine=Engine.JAX, diff_mode=diff_mode)
+        self.psr_args = psr_args
+
+    def expectation(
+        self,
+        circuit: ConvertedCircuit,
+        observable: list[ConvertedObservable] | ConvertedObservable,
+        param_values: ParamDictType = {},
+        state: ArrayLike | None = None,
+        measurement: Measurements | None = None,
+        noise: Noise | None = None,
+        mitigation: Mitigations | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> ArrayLike:
+        """Compute the expectation value of the `circuit` with the given `observable`.
+
+        Arguments:
+            circuit: A converted circuit as returned by `backend.circuit`.
+            observable: A converted observable as returned by `backend.observable`.
+            param_values: _**Already embedded**_ parameters of the circuit. See
+                [`embedding`][qadence.blocks.embedding.embedding] for more info.
+            state: Initial state.
+            measurement: Optional measurement protocol. If None, use
+                exact expectation value with a statevector simulator.
+            noise: A noise model to use.
+            mitigation: The error mitigation to use.
+            endianness: Endianness of the resulting bit strings.
+        """
+        observable = observable if isinstance(observable, list) else [observable]
+
+        if self.diff_mode == DiffMode.AD:
+            expectation = self.backend.expectation(circuit, observable, param_values, state)
+        else:
+            expectation = DifferentiableExpectation(
+                backend=self.backend,
+                circuit=circuit,
+                observable=observable,
+                param_values=param_values,
+                state=state,
+                measurement=measurement,
+                noise=noise,
+                mitigation=mitigation,
+                endianness=endianness,
+            ).psr()
+        return expectation

qadence/engines/jax/differentiable_expectation.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Tuple
+
+import jax.numpy as jnp
+from jax import Array, custom_vjp, vmap
+
+from qadence.backend import Backend as QuantumBackend
+from qadence.backend import ConvertedCircuit, ConvertedObservable
+from qadence.backends.jax_utils import (
+    tensor_to_jnp,
+)
+from qadence.blocks.utils import uuid_to_eigen
+from qadence.measurements import Measurements
+from qadence.mitigations import Mitigations
+from qadence.noise import Noise
+from qadence.types import Endianness, Engine, ParamDictType
+
+
+def compute_single_gap(eigen_vals: Array, default_val: float = 2.0) -> Array:
+    eigen_vals = eigen_vals.reshape(1, 2)
+    gaps = jnp.abs(jnp.tril(eigen_vals.T - eigen_vals))
+    return jnp.unique(jnp.where(gaps > 0.0, gaps, default_val), size=1)
+
+
+@dataclass
+class DifferentiableExpectation:
+    """A handler for differentiating expectation estimation using various engines."""
+
+    backend: QuantumBackend
+    circuit: ConvertedCircuit
+    observable: list[ConvertedObservable] | ConvertedObservable
+    param_values: ParamDictType
+    state: Array | None = None
+    measurement: Measurements | None = None
+    noise: Noise | None = None
+    mitigation: Mitigations | None = None
+    endianness: Endianness = Endianness.BIG
+    engine: Engine = Engine.JAX
+
+    def psr(self) -> Any:
+        n_obs = len(self.observable)
+
+        def expectation_fn(state: Array, values: ParamDictType, psr_params: ParamDictType) -> Array:
+            return self.backend.expectation(
+                circuit=self.circuit, observable=self.observable, param_values=values, state=state
+            )
+
+        @custom_vjp
+        def expectation(state: Array, values: ParamDictType, psr_params: ParamDictType) -> Array:
+            return expectation_fn(state, values, psr_params)
+
+        uuid_to_eigs = {
+            k: tensor_to_jnp(v) for k, v in uuid_to_eigen(self.circuit.abstract.block).items()
+        }
+        self.psr_params = {
+            k: self.param_values[k] for k in uuid_to_eigs.keys()
+        }  # Subset of params on which to perform PSR.
+
+        def expectation_fwd(state: Array, values: ParamDictType, psr_params: ParamDictType) -> Any:
+            return expectation_fn(state, values, psr_params), (
+                state,
+                values,
+                psr_params,
+            )
+
+        def expectation_bwd(res: Tuple[Array, ParamDictType, ParamDictType], tangent: Array) -> Any:
+            state, values, psr_params = res
+            grads = {}
+            # FIXME Hardcoding the single spectral_gap to 2.
+            spectral_gap = 2.0
+            shift = jnp.pi / 2
+
+            def shift_circ(param_name: str, values: dict) -> Array:
+                shifted_values = values.copy()
+                shiftvals = jnp.array(
+                    [shifted_values[param_name] + shift, shifted_values[param_name] - shift]
+                )
+
+                def _expectation(val: Array) -> Array:
+                    shifted_values[param_name] = val
+                    return expectation(state, shifted_values, psr_params)
+
+                return vmap(_expectation, in_axes=(0,))(shiftvals)
+
+            for param_name, _ in psr_params.items():
+                f_plus, f_min = shift_circ(param_name, values)
+                grad = spectral_gap * (f_plus - f_min) / (4.0 * jnp.sin(spectral_gap * shift / 2.0))
+                grads[param_name] = jnp.sum(tangent * grad, axis=1) if n_obs > 1 else tangent * grad
+            return None, None, grads
+
+        expectation.defvjp(expectation_fwd, expectation_bwd)
+        return expectation(self.state, self.param_values, self.psr_params)

qadence/engines/torch/__init__.py

@@ -0,0 +1,4 @@
+from __future__ import annotations
+
+from .differentiable_backend import DifferentiableBackend
+from .differentiable_expectation import DifferentiableExpectation

qadence/engines/torch/differentiable_backend.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from functools import partial
+
+from qadence.backend import Backend as QuantumBackend
+from qadence.backend import ConvertedCircuit, ConvertedObservable
+from qadence.engines.differentiable_backend import (
+    DifferentiableBackend as DifferentiableBackendInterface,
+)
+from qadence.engines.torch.differentiable_expectation import DifferentiableExpectation
+from qadence.extensions import get_gpsr_fns
+from qadence.measurements import Measurements
+from qadence.mitigations import Mitigations
+from qadence.noise import Noise
+from qadence.types import ArrayLike, DiffMode, Endianness, Engine, ParamDictType
+
+
+class DifferentiableBackend(DifferentiableBackendInterface):
+    """A class which wraps a QuantumBackend with the automatic differentation engine TORCH.
+
+    Arguments:
+        backend: An instance of the QuantumBackend type perform execution.
+        diff_mode: A differentiable mode supported by the differentiation engine.
+        **psr_args: Arguments that will be passed on to `DifferentiableExpectation`.
+    """
+
+    def __init__(
+        self,
+        backend: QuantumBackend,
+        diff_mode: DiffMode = DiffMode.AD,
+        **psr_args: int | float | None,
+    ) -> None:
+        super().__init__(backend=backend, engine=Engine.TORCH, diff_mode=diff_mode)
+        self.psr_args = psr_args
+
+    def expectation(
+        self,
+        circuit: ConvertedCircuit,
+        observable: list[ConvertedObservable] | ConvertedObservable,
+        param_values: ParamDictType = {},
+        state: ArrayLike | None = None,
+        measurement: Measurements | None = None,
+        noise: Noise | None = None,
+        mitigation: Mitigations | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> ArrayLike:
+        """Compute the expectation value of the `circuit` with the given `observable`.
+
+        Arguments:
+            circuit: A converted circuit as returned by `backend.circuit`.
+            observable: A converted observable as returned by `backend.observable`.
+            param_values: _**Already embedded**_ parameters of the circuit. See
+                [`embedding`][qadence.blocks.embedding.embedding] for more info.
+            state: Initial state.
+            measurement: Optional measurement protocol. If None, use
+                exact expectation value with a statevector simulator.
+            noise: A noise model to use.
+            mitigation: The error mitigation to use.
+            endianness: Endianness of the resulting bit strings.
+        """
+        observable = observable if isinstance(observable, list) else [observable]
+        differentiable_expectation = DifferentiableExpectation(
+            backend=self.backend,
+            circuit=circuit,
+            observable=observable,
+            param_values=param_values,
+            state=state,
+            measurement=measurement,
+            noise=noise,
+            mitigation=mitigation,
+            endianness=endianness,
+        )
+
+        if self.diff_mode == DiffMode.AD:
+            expectation = differentiable_expectation.ad
+        elif self.diff_mode == DiffMode.ADJOINT:
+            expectation = differentiable_expectation.adjoint
+        else:
+            try:
+                fns = get_gpsr_fns()
+                psr_fn = fns[self.diff_mode]
+            except KeyError:
+                raise ValueError(f"{self.diff_mode} differentiation mode is not supported")
+            expectation = partial(differentiable_expectation.psr, psr_fn=psr_fn, **self.psr_args)
+        return expectation()

qadence/{backends/pytorch_wrapper.py → engines/torch/differentiable_expectation.py}

@@ -13,7 +13,7 @@ from torch.nn import Module
 from qadence.backend import Backend as QuantumBackend
 from qadence.backend import Converted, ConvertedCircuit, ConvertedObservable
 from qadence.backends.adjoint import AdjointExpectation
-from qadence.backends.utils import is_pyq_shape, param_dict, pyqify, validate_state
+from qadence.backends.utils import infer_batchsize, is_pyq_shape, param_dict, pyqify, validate_state
 from qadence.blocks.abstract import AbstractBlock
 from qadence.blocks.primitive import PrimitiveBlock
 from qadence.blocks.utils import uuid_to_block, uuid_to_eigen
@@ -24,7 +24,6 @@ from qadence.mitigations import Mitigations
 from qadence.ml_tools import promote_to_tensor
 from qadence.noise import Noise
 from qadence.types import DiffMode, Endianness
-from qadence.utils import infer_batchsize
 
 
 class PSRExpectation(Function):

qadence/extensions.py

@@ -6,15 +6,30 @@ from string import Template
 from qadence.backend import Backend
 from qadence.blocks.abstract import TAbstractBlock
 from qadence.logger import get_logger
-from qadence.types import BackendName, DiffMode
+from qadence.types import BackendName, DiffMode, Engine
 
 backends_namespace = Template("qadence.backends.$name")
 
 logger = get_logger(__name__)
 
 
+def _available_engines() -> dict:
+    """Returns a dictionary of currently installed, native qadence engines."""
+    res = {}
+    for engine in Engine.list():
+        module_path = f"qadence.engines.{engine}.differentiable_backend"
+        try:
+            module = importlib.import_module(module_path)
+            DifferentiableBackendCls = getattr(module, "DifferentiableBackend")
+            res[engine] = DifferentiableBackendCls
+        except (ImportError, ModuleNotFoundError):
+            pass
+    logger.info(f"Found engines: {res.keys()}")
+    return res
+
+
 def _available_backends() -> dict:
-    """Fallback function for native Qadence available backends if extensions is not present."""
+    """Returns a dictionary of currently installed, native qadence backends."""
     res = {}
     for backend in BackendName.list():
         module_path = f"qadence.backends.{backend}.backend"
@@ -24,11 +39,12 @@ def _available_backends() -> dict:
             res[backend] = BackendCls
         except (ImportError, ModuleNotFoundError):
             pass
+    logger.info(f"Found backends: {res.keys()}")
     return res
 
 
 def _supported_gates(name: BackendName | str) -> list[TAbstractBlock]:
-    """Fallback function for native Qadence backend supported gates if extensions is not present."""
+    """Returns a list of supported gates for the queried backend 'name'."""
     from qadence import operations
 
     name = str(BackendName(name).name.lower())
@@ -102,6 +118,7 @@ try:
     set_backend_config = getattr(module, "set_backend_config")
 except ModuleNotFoundError:
     available_backends = _available_backends
+    available_engines = _available_engines
     supported_gates = _supported_gates
     get_gpsr_fns = _gpsr_fns
     set_backend_config = _set_backend_config

qadence/ml_tools/models.py

@@ -178,9 +178,16 @@ class TransformedModule(torch.nn.Module):
         if isinstance(self.model, (QuantumModel, QNN)):
             if not isinstance(x, dict):
                 x = self._format_to_dict(x)
-            return {
-                key: self._input_scaling * (val + self._input_shifting) for key, val in x.items()
-            }
+            if self.in_features == 1:
+                return {
+                    key: self._input_scaling * (val + self._input_shifting)
+                    for key, val in x.items()
+                }
+            else:
+                return {
+                    key: self._input_scaling[idx] * (val + self._input_shifting[idx])
+                    for idx, (key, val) in enumerate(x.items())
+                }
 
         else:
             assert isinstance(self.model, torch.nn.Module) and isinstance(x, Tensor)

qadence/models/quantum_model.py

@@ -18,11 +18,14 @@ from qadence.backend import (
 )
 from qadence.backends.api import backend_factory, config_factory
 from qadence.blocks.abstract import AbstractBlock
+from qadence.blocks.utils import chain, unique_parameters
 from qadence.circuit import QuantumCircuit
+from qadence.engines.differentiable_backend import DifferentiableBackend
 from qadence.logger import get_logger
 from qadence.measurements import Measurements
 from qadence.mitigations import Mitigations
 from qadence.noise import Noise
+from qadence.parameters import Parameter
 from qadence.types import DiffMode, Endianness
 
 logger = get_logger(__name__)
@@ -36,7 +39,7 @@ class QuantumModel(nn.Module):
     [here](/advanced_tutorials/custom-models.md).
     """
 
-    backend: Backend
+    backend: Backend | DifferentiableBackend
     embedding_fn: Callable
     _params: nn.ParameterDict
     _circuit: ConvertedCircuit
@@ -77,7 +80,6 @@ class QuantumModel(nn.Module):
                 f"The circuit should be of type '<class QuantumCircuit>'. Got {type(circuit)}."
             )
 
-        self.inputs = [p for p in circuit.unique_parameters if not p.trainable and not p.is_number]
         if diff_mode is None:
             raise ValueError("`diff_mode` cannot be `None` in a `QuantumModel`.")
 
@@ -90,6 +92,15 @@ class QuantumModel(nn.Module):
         else:
             observable = [observable]
 
+        def _is_feature_param(p: Parameter) -> bool:
+            return not p.trainable and not p.is_number
+
+        if observable is None:
+            self.inputs = list(filter(_is_feature_param, circuit.unique_parameters))
+        else:
+            uparams = unique_parameters(chain(circuit.block, *observable))
+            self.inputs = list(filter(_is_feature_param, uparams))
+
         conv = self.backend.convert(circuit, observable)
         self.embedding_fn = conv.embedding_fn
         self._circuit = conv.circuit

qadence/parameters.py

@@ -9,11 +9,11 @@ import sympy
 from sympy import *
 from sympy import Array, Basic, Expr, Symbol, sympify
 from sympy.physics.quantum.dagger import Dagger
-from sympytorch import SymPyModule
+from sympytorch import SymPyModule as torchSympyModule
 from torch import Tensor, heaviside, no_grad, rand, tensor
 
 from qadence.logger import get_logger
-from qadence.types import TNumber
+from qadence.types import DifferentiableExpression, Engine, TNumber
 
 # Modules to be automatically added to the qadence namespace
 __all__ = ["FeatureParameter", "Parameter", "VariationalParameter"]
@@ -190,23 +190,26 @@ def extract_original_param_entry(
     return param if not param.is_number else evaluate(param)
 
 
-def torchify(expr: Expr) -> SymPyModule:
-    """
-    Arguments:
+def heaviside_func(x: Tensor, _: Any) -> Tensor:
+    with no_grad():
+        res = heaviside(x, tensor(0.5))
+    return res
 
-        expr: An expression consisting of Parameters.
 
-    Returns:
-        A torchified, differentiable Expression.
-    """
+def torchify(expr: Expr) -> torchSympyModule:
+    extra_funcs = {sympy.core.numbers.ImaginaryUnit: 1.0j, sympy.Heaviside: heaviside_func}
+    return torchSympyModule(expressions=[sympy.N(expr)], extra_funcs=extra_funcs)
 
-    def heaviside_func(x: Tensor, _: Any) -> Tensor:
-        with no_grad():
-            res = heaviside(x, tensor(0.5))
-        return res
 
-    extra_funcs = {sympy.core.numbers.ImaginaryUnit: 1.0j, sympy.Heaviside: heaviside_func}
-    return SymPyModule(expressions=[sympy.N(expr)], extra_funcs=extra_funcs)
+def make_differentiable(expr: Expr, engine: Engine = Engine.TORCH) -> DifferentiableExpression:
+    diff_expr: DifferentiableExpression
+    if engine == Engine.JAX:
+        from qadence.backends.jax_utils import jaxify
+
+        diff_expr = jaxify(expr)
+    else:
+        diff_expr = torchify(expr)
+    return diff_expr
 
 
 def sympy_to_numeric(expr: Basic) -> TNumber:
@@ -261,7 +264,7 @@ def evaluate(expr: Expr, values: dict = {}, as_torch: bool = False) -> TNumber |
                 else:
                     raise ValueError(f"No value provided for symbol {s.name}")
         if as_torch:
-            res_value = torchify(expr)(**{s.name: tensor(v) for s, v in query.items()})
+            res_value = make_differentiable(expr)(**{s.name: tensor(v) for s, v in query.items()})
         else:
             res = expr.subs(query)
             res_value = sympy_to_numeric(res)

qadence/types.py

@@ -2,10 +2,11 @@ from __future__ import annotations
 
 import importlib
 from enum import Enum
-from typing import Iterable, Tuple, Union
+from typing import Callable, Iterable, Tuple, Union
 
 import numpy as np
 import sympy
+from numpy.typing import ArrayLike
 from torch import Tensor, pi
 
 TNumber = Union[int, float, complex]
@@ -197,6 +198,8 @@ class _BackendName(StrEnum):
     """The Braket backend."""
     PULSER = "pulser"
     """The Pulser backend."""
+    HORQRUX = "horqrux"
+    """The horqrux backend."""
 
 
 # If proprietary qadence_extensions is available, import the
@@ -386,3 +389,12 @@ class OpName(StrEnum):
 class ReadOutOptimization(StrEnum):
     MLE = "mle"
     CONSTRAINED = "constrained"
+
+
+class Engine(StrEnum):
+    TORCH = "torch"
+    JAX = "jax"
+
+
+ParamDictType = dict[str, ArrayLike]
+DifferentiableExpression = Callable[..., ArrayLike]