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

qadence/models/qnn.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
+from collections import Counter
 from typing import Callable
 
+import sympy
 from torch import Tensor
 
-from qadence.backend import BackendConfiguration
+from qadence.backend import BackendConfiguration, ConvertedObservable
 from qadence.blocks.abstract import AbstractBlock
 from qadence.circuit import QuantumCircuit
 from qadence.measurements import Measurements
+from qadence.mitigations import Mitigations
 from qadence.models.quantum_model import QuantumModel
 from qadence.noise import Noise
 from qadence.types import BackendName, DiffMode, Endianness
@@ -19,22 +22,25 @@ class QNN(QuantumModel):
     Examples:
     ```python exec="on" source="material-block" result="json"
     import torch
-    from qadence import QuantumCircuit, QNN
-    from qadence import hea, feature_map, hamiltonian_factory, Z
+    from qadence import QuantumCircuit, QNN, Z
+    from qadence import hea, feature_map, hamiltonian_factory, kron
 
     # create the circuit
     n_qubits, depth = 2, 4
-    fm = feature_map(n_qubits)
+    fm = kron(
+        feature_map(1, support=(0,), param="x"),
+        feature_map(1, support=(1,), param="y")
+    )
     ansatz = hea(n_qubits=n_qubits, depth=depth)
     circuit = QuantumCircuit(n_qubits, fm, ansatz)
-    obs_base = hamiltonian_factory(n_qubits, detuning = Z)
+    obs_base = hamiltonian_factory(n_qubits, detuning=Z)
 
     # the QNN will yield two outputs
     obs = [2.0 * obs_base, 4.0 * obs_base]
 
     # initialize and use the model
-    qnn = QNN(circuit, obs, diff_mode="ad", backend="pyqtorch")
-    y = qnn.expectation({"phi": torch.rand(3)})
+    qnn = QNN(circuit, obs, inputs=["x", "y"])
+    y = qnn(torch.rand(3, 2))
     print(str(y)) # markdown-exec: hide
     ```
     """
@@ -49,6 +55,7 @@ class QNN(QuantumModel):
         measurement: Measurements | None = None,
         noise: Noise | None = None,
         configuration: BackendConfiguration | dict | None = None,
+        inputs: list[sympy.Basic | str] | None = None,
     ):
         """Initialize the QNN.
 
@@ -59,6 +66,9 @@ class QNN(QuantumModel):
         Args:
             circuit: The quantum circuit to use for the QNN.
             transform: A transformation applied to the output of the QNN.
+            inputs: Tuple that indicates the order of variables of the tensors that are passed
+                to the model. Given input tensors `xs = torch.rand(batch_size, input_size:=2)` a QNN
+                with `inputs=("t", "x")` will assign `t, x = xs[:,0], xs[:,1]`.
             backend: The chosen quantum backend.
             diff_mode: The differentiation engine to use. Choices 'gpsr' or 'ad'.
             measurement: optional measurement protocol. If None,
@@ -67,7 +77,7 @@ class QNN(QuantumModel):
             configuration: optional configuration for the backend
         """
         super().__init__(
-            circuit=circuit,
+            circuit,
             observable=observable,
             backend=backend,
             diff_mode=diff_mode,
@@ -75,12 +85,33 @@ class QNN(QuantumModel):
             configuration=configuration,
             noise=noise,
         )
-
         if self.out_features is None:
             raise ValueError("You need to provide at least one observable in the QNN constructor")
-
         self.transform = transform if transform else lambda x: x
 
+        if (inputs is not None) and (len(self.inputs) == len(inputs)):
+            self.inputs = [sympy.symbols(x) if isinstance(x, str) else x for x in inputs]  # type: ignore[union-attr]
+        elif (inputs is None) and len(self.inputs) <= 1:
+            self.inputs = [sympy.symbols(x) if isinstance(x, str) else x for x in self.inputs]  # type: ignore[union-attr]
+        else:
+            raise ValueError(
+                """
+                Your QNN has more than one input. Please provide a list of inputs in the order of
+                your tensor domain. For example, if you want to pass
+                `xs = torch.rand(batch_size, input_size:=3)` to you QNN, where
+                ```
+                t = x[:,0]
+                x = x[:,1]
+                y = x[:,2]
+                ```
+                you have to specify
+                ```
+                QNN(circuit, observable, inputs=["t", "x", "y"])
+                ```
+                You can also pass a list of sympy symbols.
+            """
+            )
+
     def forward(
         self,
         values: dict[str, Tensor] | Tensor = None,
@@ -103,7 +134,7 @@ class QNN(QuantumModel):
         is instead `n_batches x n_observables`
 
         Args:
-            values (dict[str, Tensor] | Tensor): the values of the feature parameters
+            values: the values of the feature parameters
             state: Initial state.
             measurement: optional measurement protocol. If None,
                 use exact expectation value with a statevector simulator
@@ -114,18 +145,55 @@ class QNN(QuantumModel):
             Tensor: a tensor with the expectation value of the observables passed
                 in the constructor of the model
         """
+        return self.expectation(
+            values, state=state, measurement=measurement, noise=noise, endianness=endianness
+        )
+
+    def run(
+        self,
+        values: Tensor | dict[str, Tensor] = None,
+        state: Tensor | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> Tensor:
+        return super().run(values=self._format_to_dict(values), state=state, endianness=endianness)
+
+    def sample(
+        self,
+        values: Tensor | dict[str, Tensor] = {},
+        n_shots: int = 1000,
+        state: Tensor | None = None,
+        noise: Noise | None = None,
+        mitigation: Mitigations | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> list[Counter]:
+        return super().sample(
+            values=self._format_to_dict(values),
+            n_shots=n_shots,
+            state=state,
+            noise=noise,
+            mitigation=mitigation,
+            endianness=endianness,
+        )
+
+    def expectation(
+        self,
+        values: Tensor | dict[str, Tensor] = {},
+        observable: list[ConvertedObservable] | ConvertedObservable | None = None,
+        state: Tensor | None = None,
+        measurement: Measurements | None = None,
+        noise: Noise | None = None,
+        mitigation: Mitigations | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> Tensor:
         if values is None:
             values = {}
-        if not isinstance(values, dict):
-            values = self._format_to_dict(values)
         if measurement is None:
             measurement = self._measurement
         if noise is None:
             noise = self._noise
-
         return self.transform(
-            self.expectation(
-                values=values,
+            super().expectation(
+                values=self._format_to_dict(values),
                 state=state,
                 measurement=measurement,
                 endianness=endianness,
@@ -139,6 +207,9 @@ class QNN(QuantumModel):
         The tensor is assumed to have dimensions: n_batches x in_features where in_features
         corresponds to the number of input features of the QNN
         """
+        # for backwards compat...
+        if isinstance(values, dict):
+            return values
 
         if len(values.size()) == 1:
             values = values.reshape(-1, 1)
@@ -146,10 +217,4 @@ class QNN(QuantumModel):
         assert len(values.size()) == 2, msg
         assert values.size()[1] == self.in_features, msg
 
-        names = [p.name for p in self.inputs]
-        res = {}
-        for i, name in enumerate(names):
-            res[name] = values[:, i]
-        return res
-
-    # TODO: Implement derivatives w.r.t. to inputs
+        return {var.name: values[:, self.inputs.index(var)] for var in self.inputs}

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/operations.py

@@ -31,6 +31,7 @@ from qadence.blocks.analog import (
     WaitBlock,
 )
 from qadence.blocks.block_to_tensor import block_to_tensor
+from qadence.blocks.primitive import ProjectorBlock
 from qadence.blocks.utils import (
     add,  # noqa
     block_is_commuting_hamiltonian,
@@ -117,9 +118,6 @@ class X(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return tensor([-1, 1], dtype=cdouble)
 
-    def dagger(self) -> X:
-        return self
-
 
 class Y(PrimitiveBlock):
     """The Y gate."""
@@ -141,9 +139,6 @@ class Y(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return tensor([-1, 1], dtype=cdouble)
 
-    def dagger(self) -> Y:
-        return self
-
 
 class Z(PrimitiveBlock):
     """The Z gate."""
@@ -165,17 +160,36 @@ class Z(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return tensor([-1, 1], dtype=cdouble)
 
-    def dagger(self) -> Z:
-        return self
 
+class Projector(ProjectorBlock):
+    """The projector operator."""
+
+    name = OpName.PROJ
+
+    def __init__(
+        self,
+        ket: str,
+        bra: str,
+        qubit_support: int | tuple[int, ...],
+    ):
+        super().__init__(ket=ket, bra=bra, qubit_support=qubit_support)
+
+    @property
+    def generator(self) -> None:
+        raise ValueError("Property `generator` not available for non-unitary operator.")
+
+    @property
+    def eigenvalues_generator(self) -> None:
+        raise ValueError("Property `eigenvalues_generator` not available for non-unitary operator.")
 
-class N(PrimitiveBlock):
+
+class N(Projector):
     """The N = (1/2)(I-Z) operator."""
 
     name = OpName.N
 
-    def __init__(self, target: int):
-        super().__init__((target,))
+    def __init__(self, target: int, state: str = "1"):
+        super().__init__(ket=state, bra=state, qubit_support=(target,))
 
     @property
     def generator(self) -> None:
@@ -189,9 +203,6 @@ class N(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return tensor([0, 1], dtype=cdouble)
 
-    def dagger(self) -> N:
-        return self
-
 
 class S(PrimitiveBlock):
     """The S / Phase gate."""
@@ -296,9 +307,6 @@ class I(PrimitiveBlock):
     def __ascii__(self, console: Console) -> Padding:
         return Padding("──────", (1, 1, 1, 1))
 
-    def dagger(self) -> I:
-        return I(*self.qubit_support)
-
 
 TPauliBlock = Union[X, Y, Z, I, N]
 
@@ -320,9 +328,6 @@ class H(PrimitiveBlock):
     def eigenvalues(self) -> Tensor:
         return torch.tensor([-1, 1], dtype=cdouble)
 
-    def dagger(self) -> H:
-        return H(*self.qubit_support)
-
 
 class Zero(PrimitiveBlock):
     name = OpName.ZERO
@@ -363,9 +368,6 @@ class Zero(PrimitiveBlock):
     def __pow__(self, power: int) -> AbstractBlock:
         return self
 
-    def dagger(self) -> Zero:
-        return Zero()
-
 
 class RX(ParametricBlock):
     """The Rx gate."""
@@ -668,7 +670,7 @@ class CNOT(ControlBlock):
     name = OpName.CNOT
 
     def __init__(self, control: int, target: int) -> None:
-        self.generator = kron((I(control) - Z(control)) * 0.5, X(target) - I(target))
+        self.generator = kron(N(control), X(target) - I(target))
         super().__init__((control,), X(target))
 
     @property
@@ -691,17 +693,12 @@ class CNOT(ControlBlock):
             tree.add(self._block_title)
         return tree
 
-    def dagger(self) -> CNOT:
-        return CNOT(*self.qubit_support)
-
 
 class MCZ(ControlBlock):
     name = OpName.MCZ
 
     def __init__(self, control: tuple[int, ...], target: int) -> None:
-        self.generator = kron(
-            *[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], Z(target) - I(target)
-        )
+        self.generator = kron(*[N(qubit) for qubit in control], Z(target) - I(target))
         super().__init__(control, Z(target))
 
     @property
@@ -724,9 +721,6 @@ class MCZ(ControlBlock):
             tree.add(self._block_title)
         return tree
 
-    def dagger(self) -> MCZ:
-        return MCZ(self.qubit_support[:-1], self.qubit_support[-1])
-
 
 class CZ(MCZ):
     """The CZ gate."""
@@ -736,9 +730,6 @@ class CZ(MCZ):
     def __init__(self, control: int, target: int) -> None:
         super().__init__((control,), target)
 
-    def dagger(self) -> CZ:
-        return CZ(self.qubit_support[-2], self.qubit_support[-1])
-
 
 class MCRX(ParametricControlBlock):
     name = OpName.MCRX
@@ -749,7 +740,7 @@ class MCRX(ParametricControlBlock):
         target: int,
         parameter: Parameter | TNumber | sympy.Expr | str,
     ) -> None:
-        self.generator = kron(*[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], X(target))
+        self.generator = kron(*[N(qubit) for qubit in control], X(target))
         super().__init__(control, RX(target, parameter))
 
     @classmethod
@@ -792,7 +783,7 @@ class MCRY(ParametricControlBlock):
         target: int,
         parameter: Parameter | TNumber | sympy.Expr | str,
     ) -> None:
-        self.generator = kron(*[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], Y(target))
+        self.generator = kron(*[N(qubit) for qubit in control], Y(target))
         super().__init__(control, RY(target, parameter))
 
     @classmethod
@@ -821,7 +812,7 @@ class CRY(MCRY):
         self,
         control: int,
         target: int,
-        parameter: Parameter | TNumber | sympy.Expr | str,
+        parameter: TParameter,
     ):
         super().__init__((control,), target, parameter)
 
@@ -835,7 +826,7 @@ class MCRZ(ParametricControlBlock):
         target: int,
         parameter: Parameter | TNumber | sympy.Expr | str,
     ) -> None:
-        self.generator = kron(*[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], Z(target))
+        self.generator = kron(*[N(qubit) for qubit in control], Z(target))
         super().__init__(control, RZ(target, parameter))
 
     @classmethod
@@ -878,10 +869,10 @@ class CSWAP(ControlBlock):
         if isinstance(control, tuple):
             control = control[0]
 
-        a00m = 0.5 * (Z(control) - I(control))
-        a00p = -0.5 * (Z(control) + I(control))
-        a11 = 0.5 * (Z(target1) - I(target1))
-        a22 = -0.5 * (Z(target2) + I(target2))
+        a00m = -N(target=control)
+        a00p = -N(target=control, state="0")
+        a11 = -N(target=target1)
+        a22 = -N(target=target2, state="0")
         a12 = 0.5 * (chain(X(target1), Z(target1)) + X(target1))
         a21 = 0.5 * (chain(Z(target2), X(target2)) + X(target2))
         no_effect = kron(a00m, I(target1), I(target2))
@@ -906,9 +897,6 @@ class CSWAP(ControlBlock):
     def nqubits(self) -> int:
         return 3
 
-    def dagger(self) -> CSWAP:
-        return CSWAP(*self.qubit_support)
-
 
 class T(PrimitiveBlock):
     """The T gate."""
@@ -994,9 +982,6 @@ class SWAP(PrimitiveBlock):
         s = f"{self.name}({c}, {t})"
         return s if self.tag is None else (s + rf" \[tag: {self.tag}]")
 
-    def dagger(self) -> SWAP:
-        return SWAP(*self.qubit_support)
-
 
 class AnalogSWAP(HamEvo):
     """
@@ -1031,9 +1016,7 @@ class MCPHASE(ParametricControlBlock):
         target: int,
         parameter: Parameter | TNumber | sympy.Expr | str,
     ) -> None:
-        self.generator = kron(
-            *[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], Z(target) - I(target)
-        )
+        self.generator = kron(*[N(qubit) for qubit in control], Z(target) - I(target))
         super().__init__(control, PHASE(target, parameter))
 
     @classmethod
@@ -1082,14 +1065,9 @@ class Toffoli(ControlBlock):
     name = OpName.TOFFOLI
 
     def __init__(self, control: tuple[int, ...], target: int) -> None:
-        self.generator = kron(
-            *[(I(qubit) - Z(qubit)) * 0.5 for qubit in control], X(target) - I(target)
-        )
+        self.generator = kron(*[N(qubit) for qubit in control], X(target) - I(target))
         super().__init__(control, X(target))
 
-    def dagger(self) -> Toffoli:
-        return Toffoli(self.qubit_support[:-1], self.qubit_support[-1])
-
     @property
     def n_qubits(self) -> int:
         return len(self.qubit_support)
@@ -1129,6 +1107,7 @@ def _cast(T: Any, val: Any) -> Any:
 def wait(
     duration: TNumber | sympy.Basic,
     qubit_support: str | QubitSupport | tuple = "global",
+    add_pattern: bool = True,
 ) -> WaitBlock:
     """Constructs a [`WaitBlock`][qadence.blocks.analog.WaitBlock].
 
@@ -1142,7 +1121,7 @@ def wait(
     """
     q = _cast(QubitSupport, qubit_support)
     ps = ParamMap(duration=duration)
-    return WaitBlock(parameters=ps, qubit_support=q)
+    return WaitBlock(parameters=ps, qubit_support=q, add_pattern=add_pattern)
 
 
 def entangle(
@@ -1160,6 +1139,7 @@ def AnalogRot(
     delta: float | str | Parameter = 0,
     phase: float | str | Parameter = 0,
     qubit_support: str | QubitSupport | Tuple = "global",
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     """General analog rotation operation.
 
@@ -1174,18 +1154,20 @@ def AnalogRot(
         ConstantAnalogRotation
     """
     q = _cast(QubitSupport, qubit_support)
-    if isinstance(duration, str):
-        duration = Parameter(duration)
-    alpha = duration * sympy.sqrt(omega**2 + delta**2) / 1000  # type: ignore [operator]
-
+    duration = Parameter(duration)
+    omega = Parameter(omega)
+    delta = Parameter(delta)
+    phase = Parameter(phase)
+    alpha = duration * sympy.sqrt(omega**2 + delta**2) / 1000
     ps = ParamMap(alpha=alpha, duration=duration, omega=omega, delta=delta, phase=phase)
-    return ConstantAnalogRotation(parameters=ps, qubit_support=q)
+    return ConstantAnalogRotation(parameters=ps, qubit_support=q, add_pattern=add_pattern)
 
 
 def _analog_rot(
     angle: float | str | Parameter,
     qubit_support: str | QubitSupport | Tuple,
     phase: float,
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     q = _cast(QubitSupport, qubit_support)
     # assuming some arbitrary omega = π rad/μs
@@ -1200,12 +1182,13 @@ def _analog_rot(
     # and compute omega like this:
     # omega = alpha / duration * 1000
     ps = ParamMap(alpha=alpha, duration=duration, omega=omega, delta=0, phase=phase)
-    return ConstantAnalogRotation(parameters=ps, qubit_support=q)
+    return ConstantAnalogRotation(parameters=ps, qubit_support=q, add_pattern=add_pattern)
 
 
 def AnalogRX(
     angle: float | str | Parameter,
     qubit_support: str | QubitSupport | Tuple = "global",
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     """Analog X rotation.
 
@@ -1223,12 +1206,13 @@ def AnalogRX(
     Returns:
         ConstantAnalogRotation
     """
-    return _analog_rot(angle, qubit_support, phase=0)
+    return _analog_rot(angle, qubit_support, phase=0, add_pattern=add_pattern)
 
 
 def AnalogRY(
     angle: float | str | Parameter,
     qubit_support: str | QubitSupport | Tuple = "global",
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     """Analog Y rotation.
 
@@ -1245,12 +1229,13 @@ def AnalogRY(
     Returns:
         ConstantAnalogRotation
     """
-    return _analog_rot(angle, qubit_support, phase=-np.pi / 2)
+    return _analog_rot(angle, qubit_support, phase=-np.pi / 2, add_pattern=add_pattern)
 
 
 def AnalogRZ(
     angle: float | str | Parameter,
     qubit_support: str | QubitSupport | Tuple = "global",
+    add_pattern: bool = True,
 ) -> ConstantAnalogRotation:
     """Analog Z rotation. Shorthand for [`AnalogRot`][qadence.operations.AnalogRot]:
     ```
@@ -1263,7 +1248,7 @@ def AnalogRZ(
     delta = np.pi
     duration = alpha / delta * 1000
     ps = ParamMap(alpha=alpha, duration=duration, omega=0, delta=delta, phase=0.0)
-    return ConstantAnalogRotation(qubit_support=q, parameters=ps)
+    return ConstantAnalogRotation(qubit_support=q, parameters=ps, add_pattern=add_pattern)
 
 
 # gate sets
@@ -1288,4 +1273,4 @@ analog_gateset = [
     entangle,
     wait,
 ]
-non_unitary_gateset = [Zero, N]
+non_unitary_gateset = [Zero, N, Projector]

qadence/parameters.py

@@ -8,17 +8,19 @@ import numpy as np
 import sympy
 from sympy import *
 from sympy import Array, Basic, Expr, Symbol, sympify
-from sympytorch import SymPyModule
-from torch import Tensor, rand, tensor
+from sympy.physics.quantum.dagger import Dagger
+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"]
 
 logger = get_logger(__file__)
 
+dagger_expression = Dagger
 
 ParameterJSONSchema = {
     "$schema": "http://json-schema.org/draft-07/schema#",
@@ -188,17 +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.
-    """
-    extra_funcs = {sympy.core.numbers.ImaginaryUnit: 1.0j}
-    return SymPyModule(expressions=[sympy.N(expr)], extra_funcs=extra_funcs)
+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 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:
@@ -253,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)