qadence 1.7.4__py3-none-any.whl → 1.7.6__py3-none-any.whl

qadence/analog/addressing.py

@@ -59,9 +59,13 @@ class AddressingPattern:
     ) -> dict:
         # augment weight dict if needed
         weights = {
-            i: Parameter(0.0)
-            if i not in weights
-            else (Parameter(weights[i]) if not isinstance(weights[i], Parameter) else weights[i])
+            i: (
+                Parameter(0.0)
+                if i not in weights
+                else (
+                    Parameter(weights[i]) if not isinstance(weights[i], Parameter) else weights[i]
+                )
+            )
             for i in range(self.n_qubits)
         }
 

qadence/backends/api.py

@@ -3,6 +3,9 @@ from __future__ import annotations
 from qadence.backend import Backend, BackendConfiguration
 from qadence.engines.differentiable_backend import DifferentiableBackend
 from qadence.extensions import (
+    BackendNotFoundError,
+    ConfigNotFoundError,
+    EngineNotFoundError,
     import_backend,
     import_config,
     import_engine,
@@ -49,12 +52,9 @@ def backend_factory(
             diff_backend_cls = import_engine(backend_inst.engine)
             backend_inst = diff_backend_cls(backend=backend_inst, diff_mode=DiffMode(diff_mode))  # type: ignore[operator]
         return backend_inst
-    except Exception as e:
-        msg = f"The requested backend '{backend}' is either not installed\
-            or could not be imported due to {e}."
-        logger.error(msg)
-        raise Exception(msg)
-    # Set backend configurations which depend on the differentiation mode
+    except (BackendNotFoundError, EngineNotFoundError, ConfigNotFoundError) as e:
+        logger.error(e.msg)
+        raise e
 
 
 def config_factory(backend_name: BackendName | str, config: dict) -> BackendConfiguration:
@@ -62,6 +62,7 @@ def config_factory(backend_name: BackendName | str, config: dict) -> BackendConf
     try:
         BackendConfigCls = import_config(backend_name)
         cfg = BackendConfigCls(**config)  # type: ignore[operator]
-    except Exception as e:
-        logger.debug(f"Unable to import config for backend {backend_name} due to {e}.")
+    except ConfigNotFoundError as e:
+        logger.error(e.msg)
+        raise e
     return cfg

qadence/backends/gpsr.py

@@ -11,13 +11,29 @@ from qadence.utils import _round_complex
 
 
 def general_psr(spectrum: Tensor, n_eqs: int | None = None, shift_prefac: float = 0.5) -> Callable:
+    """Define whether single_gap_psr or multi_gap_psr is used.
+
+    Args:
+        spectrum (Tensor): Spectrum of the operation we apply PSR onto.
+        n_eqs (int | None, optional): Number of equations. Defaults to None.
+            If provided, we keep the n_eqs higher spectral gaps.
+        shift_prefac (float, optional): Shift prefactor. Defaults to 0.5.
+
+    Returns:
+        Callable: single_gap_psr or multi_gap_psr function for
+            concerned operation.
+    """
     diffs = _round_complex(spectrum - spectrum.reshape(-1, 1))
     sorted_unique_spectral_gaps = torch.unique(torch.abs(torch.tril(diffs)))
 
     # We have to filter out zeros
     sorted_unique_spectral_gaps = sorted_unique_spectral_gaps[sorted_unique_spectral_gaps > 0]
-    n_eqs = len(sorted_unique_spectral_gaps)
-    sorted_unique_spectral_gaps = torch.tensor(list(sorted_unique_spectral_gaps))
+    n_eqs = (
+        len(sorted_unique_spectral_gaps)
+        if n_eqs is None
+        else min(n_eqs, len(sorted_unique_spectral_gaps))
+    )
+    sorted_unique_spectral_gaps = torch.tensor(list(sorted_unique_spectral_gaps)[:n_eqs])
 
     if n_eqs == 1:
         return single_gap_psr

qadence/backends/horqrux/convert_ops.py

@@ -7,11 +7,11 @@ from operator import add
 from typing import Any, Callable, Dict
 
 import jax.numpy as jnp
-from horqrux.abstract import Primitive as Gate
 from horqrux.analog import _HamiltonianEvolution as NativeHorqHEvo
 from horqrux.apply import apply_gate
 from horqrux.parametric import RX, RY, RZ
 from horqrux.primitive import NOT, SWAP, H, I, X, Y, Z
+from horqrux.primitive import Primitive as Gate
 from horqrux.utils import inner
 from jax import Array
 from jax.scipy.linalg import expm

qadence/backends/pyqtorch/backend.py

@@ -71,19 +71,16 @@ class Backend(BackendInterface):
     def observable(self, observable: AbstractBlock, n_qubits: int) -> ConvertedObservable:
         # make sure only leaves, i.e. primitive blocks are scaled
         transpilations = [
-            lambda block: chain_single_qubit_ops(block)
-            if self.config.use_single_qubit_composition
-            else flatten(block),
+            lambda block: (
+                chain_single_qubit_ops(block)
+                if self.config.use_single_qubit_composition
+                else flatten(block)
+            ),
             scale_primitive_blocks_only,
         ]
         block = transpile(*transpilations)(observable)  # type: ignore[call-overload]
         operations = convert_block(block, n_qubits, self.config)
-        obs_cls = (
-            pyq.DiagonalObservable
-            if block._is_diag_pauli and not block.is_parametric
-            else pyq.Observable
-        )
-        native = obs_cls(n_qubits=n_qubits, operations=operations)
+        native = pyq.Observable(operations=operations)
         return ConvertedObservable(native=native, abstract=block, original=observable)
 
     def run(
@@ -140,7 +137,7 @@ class Backend(BackendInterface):
         )
         observable = observable if isinstance(observable, list) else [observable]
         _expectation = torch.hstack(
-            [obs.native(state, param_values).reshape(-1, 1) for obs in observable]
+            [obs.native.expectation(state, param_values).reshape(-1, 1) for obs in observable]
         )
         return _expectation
 
@@ -169,7 +166,7 @@ class Backend(BackendInterface):
         observables = observable if isinstance(observable, list) else [observable]
         for vals in to_list_of_dicts(param_values):
             wf = self.run(circuit, vals, state, endianness, pyqify_state=True, unpyqify_state=False)
-            exs = torch.cat([obs.native(wf, vals) for obs in observables], 0)
+            exs = torch.cat([obs.native.expectation(wf, vals) for obs in observables], 0)
             list_expvals.append(exs)
 
         batch_expvals = torch.vstack(list_expvals)

qadence/backends/pyqtorch/convert_ops.py

@@ -7,7 +7,6 @@ from typing import Any, Sequence, Tuple
 import pyqtorch as pyq
 import sympy
 import torch
-from pyqtorch.apply import apply_operator
 from pyqtorch.embed import Embedding
 from pyqtorch.matrices import _dagger
 from pyqtorch.time_dependent.sesolve import sesolve
@@ -45,7 +44,6 @@ from qadence.blocks import (
 )
 from qadence.blocks.block_to_tensor import (
     _block_to_tensor_embedded,
-    block_to_tensor,
 )
 from qadence.blocks.primitive import ProjectorBlock
 from qadence.blocks.utils import parameters
@@ -78,6 +76,14 @@ def is_single_qubit_chain(block: AbstractBlock) -> bool:
     )
 
 
+def extract_parameter(block: ScaleBlock | ParametricBlock, config: Configuration) -> str | Tensor:
+    return (
+        tensor([block.parameters.parameter], dtype=float64)
+        if not block.is_parametric
+        else config.get_param_name(block)[0]
+    )
+
+
 def convert_block(
     block: AbstractBlock, n_qubits: int = None, config: Configuration = None
 ) -> Sequence[Module | Tensor | str | sympy.Expr]:
@@ -94,31 +100,45 @@ def convert_block(
 
     if isinstance(block, ScaleBlock):
         scaled_ops = convert_block(block.block, n_qubits, config)
-        scale = (
-            tensor([block.parameters.parameter], dtype=float64)
-            if not block.is_parametric
-            else config.get_param_name(block)[0]
-        )
+        scale = extract_parameter(block, config)
         return [pyq.Scale(pyq.Sequence(scaled_ops), scale)]
 
     elif isinstance(block, TimeEvolutionBlock):
-        # TODO add native pyq hamevo
-        # generator = convert_block(block.generator, n_qubits, config)[0]  # type: ignore[arg-type]
-        # time_param = config.get_param_name(block)[0]
-        # is_parametric = (
-        #     block.generator.is_parametric if isinstance(block.generator, AbstractBlock) else False
-        # )
-        # return [
-        #     pyq.HamiltonianEvolution(
-        #         qubit_support=qubit_support,
-        #         generator=generator,
-        #         time=time_param,
-        #         generator_parametric=is_parametric,  # type: ignore[union-attr]
-        #     )
-        # ]
-        return [PyQHamiltonianEvolution(qubit_support, n_qubits, block, config)]
+        if getattr(block.generator, "is_time_dependent", False):
+            return [PyQTimeDependentEvolution(qubit_support, n_qubits, block, config)]
+        else:
+            if isinstance(block.generator, sympy.Basic):
+                generator = config.get_param_name(block)[1]
+            elif isinstance(block.generator, Tensor):
+                m = block.generator.to(dtype=cdouble)
+                generator = convert_block(
+                    MatrixBlock(
+                        m,
+                        qubit_support=qubit_support,
+                        check_unitary=False,
+                        check_hermitian=True,
+                    )
+                )[0]
+            else:
+                generator = convert_block(block.generator, n_qubits, config)[0]  # type: ignore[arg-type]
+            time_param = config.get_param_name(block)[0]
+            is_parametric = (
+                block.generator.is_parametric
+                if isinstance(block.generator, AbstractBlock)
+                else False
+            )
+            return [
+                pyq.HamiltonianEvolution(
+                    qubit_support=qubit_support,
+                    generator=generator,
+                    time=time_param,
+                    generator_parametric=is_parametric,  # type: ignore[union-attr]
+                    cache_length=0,
+                )
+            ]
+
     elif isinstance(block, MatrixBlock):
-        return [pyq.primitive.Primitive(block.matrix, block.qubit_support)]
+        return [pyq.primitives.Primitive(block.matrix, block.qubit_support)]
     elif isinstance(block, CompositeBlock):
         ops = list(flatten(*(convert_block(b, n_qubits, config) for b in block.blocks)))
         if isinstance(block, AddBlock):
@@ -142,14 +162,14 @@ def convert_block(
             if isinstance(block, U):
                 op = pyq_cls(qubit_support[0], *config.get_param_name(block))
             else:
-                op = pyq_cls(qubit_support[0], config.get_param_name(block)[0])
+                op = pyq_cls(qubit_support[0], extract_parameter(block, config))
         else:
             op = pyq_cls(qubit_support[0])
         return [op]
     elif isinstance(block, tuple(two_qubit_gateset)):
         pyq_cls = getattr(pyq, block.name)
         if isinstance(block, ParametricBlock):
-            op = pyq_cls(qubit_support[0], qubit_support[1], config.get_param_name(block)[0])
+            op = pyq_cls(qubit_support[0], qubit_support[1], extract_parameter(block, config))
         else:
             op = pyq_cls(qubit_support[0], qubit_support[1])
         return [op]
@@ -157,7 +177,7 @@ def convert_block(
         block_name = block.name[1:] if block.name.startswith("M") else block.name
         pyq_cls = getattr(pyq, block_name)
         if isinstance(block, ParametricBlock):
-            op = pyq_cls(qubit_support[:-1], qubit_support[-1], config.get_param_name(block)[0])
+            op = pyq_cls(qubit_support[:-1], qubit_support[-1], extract_parameter(block, config))
         else:
             if "CSWAP" in block_name:
                 op = pyq_cls(qubit_support[:-2], qubit_support[-2:])
@@ -172,7 +192,7 @@ def convert_block(
         )
 
 
-class PyQHamiltonianEvolution(Module):
+class PyQTimeDependentEvolution(Module):
     def __init__(
         self,
         qubit_support: Tuple[int, ...],
@@ -188,50 +208,17 @@ class PyQHamiltonianEvolution(Module):
         self.hmat: Tensor
         self.config = config
 
-        if isinstance(block.generator, AbstractBlock) and not block.generator.is_parametric:
-            hmat = block_to_tensor(
-                block.generator,
-                qubit_support=self.qubit_support,
-                use_full_support=False,
-            )
-            hmat = hmat.permute(1, 2, 0)
-            self.register_buffer("hmat", hmat)
-            self._hamiltonian = lambda self, values: self.hmat
-
-        elif isinstance(block.generator, Tensor):
-            m = block.generator.to(dtype=cdouble)
-            hmat = block_to_tensor(
-                MatrixBlock(
-                    m,
-                    qubit_support=block.qubit_support,
-                    check_unitary=False,
-                    check_hermitian=True,
-                ),
+        def _hamiltonian(self: PyQTimeDependentEvolution, values: dict[str, Tensor]) -> Tensor:
+            hmat = _block_to_tensor_embedded(
+                block.generator,  # type: ignore[arg-type]
+                values=values,
                 qubit_support=self.qubit_support,
                 use_full_support=False,
+                device=self.device,
             )
-            hmat = hmat.permute(1, 2, 0)
-            self.register_buffer("hmat", hmat)
-            self._hamiltonian = lambda self, values: self.hmat
-
-        elif isinstance(block.generator, sympy.Basic):
-            self._hamiltonian = (
-                lambda self, values: values[self.param_names[1]].squeeze(3).permute(1, 2, 0)
-            )
-            # FIXME Why are we squeezing
-        else:
-
-            def _hamiltonian(self: PyQHamiltonianEvolution, values: dict[str, Tensor]) -> Tensor:
-                hmat = _block_to_tensor_embedded(
-                    block.generator,  # type: ignore[arg-type]
-                    values=values,
-                    qubit_support=self.qubit_support,
-                    use_full_support=False,
-                    device=self.device,
-                )
-                return hmat.permute(1, 2, 0)
+            return hmat.permute(1, 2, 0)
 
-            self._hamiltonian = _hamiltonian
+        self._hamiltonian = _hamiltonian
 
         self._time_evolution = lambda values: values[self.param_names[0]]
         self._device: torch_device = (
@@ -322,61 +309,51 @@ class PyQHamiltonianEvolution(Module):
         values: dict[str, Tensor] | ParameterDict = dict(),
         embedding: Embedding | None = None,
     ) -> Tensor:
-        if getattr(self.block.generator, "is_time_dependent", False):  # type: ignore [union-attr]
-
-            def Ht(t: Tensor | float) -> Tensor:
-                # values dict has to change with new value of t
-                # initial value of a feature parameter inside generator block
-                # has to be inferred here
-                new_vals = dict()
-                for str_expr, val in values.items():
-                    expr = sympy.sympify(str_expr)
-                    t_symb = sympy.Symbol(self._get_time_parameter())
-                    free_symbols = expr.free_symbols
-                    if t_symb in free_symbols:
-                        # create substitution list for time and feature params
-                        subs_list = [(t_symb, t)]
-
-                        if len(free_symbols) > 1:
-                            # get feature param symbols
-                            feat_symbols = free_symbols.difference(set([t_symb]))
-
-                            # get feature param values
-                            feat_vals = values["orig_param_values"]
-
-                            # update substitution list with feature param values
-                            for fs in feat_symbols:
-                                subs_list.append((fs, feat_vals[str(fs)]))
-
-                        # evaluate expression with new time param value
-                        new_vals[str_expr] = torch.tensor(float(expr.subs(subs_list)))
-                    else:
-                        # expression doesn't contain time parameter - copy it as is
-                        new_vals[str_expr] = val
-
-                # get matrix form of generator
-                hmat = _block_to_tensor_embedded(
-                    self.block.generator,  # type: ignore[arg-type]
-                    values=new_vals,
-                    qubit_support=self.qubit_support,
-                    use_full_support=False,
-                    device=self.device,
-                ).squeeze(0)
-
-                return hmat
-
-            tsave = torch.linspace(0, self.block.duration, self.config.n_steps_hevo)  # type: ignore [attr-defined]
-            result = pyqify(
-                sesolve(Ht, unpyqify(state).T[:, 0:1], tsave, self.config.ode_solver).states[-1].T
-            )
-        else:
-            result = apply_operator(
-                state,
-                self.unitary(values),
-                self.qubit_support,
-                self.n_qubits,
-                self.batch_size,
-            )
+        def Ht(t: Tensor | float) -> Tensor:
+            # values dict has to change with new value of t
+            # initial value of a feature parameter inside generator block
+            # has to be inferred here
+            new_vals = dict()
+            for str_expr, val in values.items():
+                expr = sympy.sympify(str_expr)
+                t_symb = sympy.Symbol(self._get_time_parameter())
+                free_symbols = expr.free_symbols
+                if t_symb in free_symbols:
+                    # create substitution list for time and feature params
+                    subs_list = [(t_symb, t)]
+
+                    if len(free_symbols) > 1:
+                        # get feature param symbols
+                        feat_symbols = free_symbols.difference(set([t_symb]))
+
+                        # get feature param values
+                        feat_vals = values["orig_param_values"]
+
+                        # update substitution list with feature param values
+                        for fs in feat_symbols:
+                            subs_list.append((fs, feat_vals[str(fs)]))
+
+                    # evaluate expression with new time param value
+                    new_vals[str_expr] = torch.tensor(float(expr.subs(subs_list)))
+                else:
+                    # expression doesn't contain time parameter - copy it as is
+                    new_vals[str_expr] = val
+
+            # get matrix form of generator
+            hmat = _block_to_tensor_embedded(
+                self.block.generator,  # type: ignore[arg-type]
+                values=new_vals,
+                qubit_support=self.qubit_support,
+                use_full_support=False,
+                device=self.device,
+            ).squeeze(0)
+
+            return hmat
+
+        tsave = torch.linspace(0, self.block.duration, self.config.n_steps_hevo)  # type: ignore [attr-defined]
+        result = pyqify(
+            sesolve(Ht, unpyqify(state).T[:, 0:1], tsave, self.config.ode_solver).states[-1].T
+        )
 
         return result
 
@@ -388,7 +365,7 @@ class PyQHamiltonianEvolution(Module):
     def dtype(self) -> torch_dtype:
         return self._dtype
 
-    def to(self, *args: Any, **kwargs: Any) -> PyQHamiltonianEvolution:
+    def to(self, *args: Any, **kwargs: Any) -> PyQTimeDependentEvolution:
         if hasattr(self, "hmat"):
             self.hmat = self.hmat.to(*args, **kwargs)
             self._device = self.hmat.device

qadence/backends/utils.py

@@ -9,7 +9,7 @@ import pyqtorch as pyq
 import torch
 from numpy.typing import ArrayLike
 from pyqtorch.apply import apply_operator
-from pyqtorch.parametric import Parametric as PyQParametric
+from pyqtorch.primitives import Parametric as PyQParametric
 from torch import (
     Tensor,
     cat,

qadence/blocks/composite.py

@@ -129,9 +129,11 @@ class CompositeBlock(AbstractBlock):
         from qadence.blocks.utils import _construct, tag
 
         blocks = [
-            getattr(operations, b["type"])._from_dict(b)
-            if hasattr(operations, b["type"])
-            else getattr(qadenceblocks, b["type"])._from_dict(b)
+            (
+                getattr(operations, b["type"])._from_dict(b)
+                if hasattr(operations, b["type"])
+                else getattr(qadenceblocks, b["type"])._from_dict(b)
+            )
             for b in d["blocks"]
         ]
         block = _construct(cls, blocks)  # type: ignore[arg-type]

qadence/blocks/utils.py

@@ -263,11 +263,29 @@ def expression_to_uuids(block: AbstractBlock) -> dict[Expr, list[str]]:
     return expr_to_uuid
 
 
-def uuid_to_eigen(block: AbstractBlock) -> dict[str, Tensor]:
+def uuid_to_eigen(
+    block: AbstractBlock, rescale_eigenvals_timeevo: bool = False
+) -> dict[str, Tensor]:
     """Creates a mapping between a parametric block's param_id and its' eigenvalues.
 
     This method is needed for constructing the PSR rules for a given block.
 
+    A PSR shift factor is also added in the mapping for dealing
+    with the time evolution case as it requires rescaling.
+
+    Args:
+        block (AbstractBlock): Block input
+        rescale_eigenvals_timeevo (bool, optional): If True, rescale
+        eigenvalues and shift factor
+        by 2 times spectral gap
+        for the TimeEvolutionBlock case to allow
+        differientiating with Hamevo.
+        Defaults to False.
+
+    Returns:
+        dict[str, Tensor]: Mapping between block's param_id, eigenvalues and
+        PSR shift.
+
     !!! warn
         Will ignore eigenvalues of AnalogBlocks that are not yet computed.
     """
@@ -276,7 +294,23 @@ def uuid_to_eigen(block: AbstractBlock) -> dict[str, Tensor]:
     for uuid, b in uuid_to_block(block).items():
         if b.eigenvalues_generator is not None:
             if b.eigenvalues_generator.numel() > 0:
-                result[uuid] = b.eigenvalues_generator
+                # GPSR assumes a factor 0.5 for differentiation
+                # so need rescaling
+                if isinstance(b, TimeEvolutionBlock) and rescale_eigenvals_timeevo:
+                    if b.eigenvalues_generator.numel() > 1:
+                        result[uuid] = (
+                            b.eigenvalues_generator * 2.0,
+                            0.5,
+                        )
+                    else:
+                        result[uuid] = (
+                            b.eigenvalues_generator * 2.0,
+                            1.0 / (b.eigenvalues_generator.item() * 2.0)
+                            if len(b.eigenvalues_generator) == 1
+                            else 1.0,
+                        )
+                else:
+                    result[uuid] = (b.eigenvalues_generator, 1.0)
 
                 # leave only angle parameter uuid with eigenvals for ConstantAnalogRotation block
                 if isinstance(block, ConstantAnalogRotation):

qadence/constructors/utils.py

@@ -7,61 +7,59 @@ import sympy
 from qadence.blocks import KronBlock, kron
 from qadence.operations import RY
 from qadence.parameters import FeatureParameter, Parameter
-from qadence.types import PI
+from qadence.types import PI, BasisSet, MultivariateStrategy, ReuploadScaling
 
 
-def generator_prefactor(spectrum: str, qubit_index: int) -> float | int:
+def generator_prefactor(reupload_scaling: ReuploadScaling, qubit_index: int) -> float | int:
     """Converts a spectrum string, e.g. tower or exponential.
 
     The result is the correct generator prefactor.
     """
-    spectrum = spectrum.lower()
     conversion_dict: dict[str, float | int] = {
-        "simple": 1,
-        "tower": qubit_index + 1,
-        "exponential": 2 * PI / (2 ** (qubit_index + 1)),
+        ReuploadScaling.CONSTANT: 1,
+        ReuploadScaling.TOWER: qubit_index + 1,
+        ReuploadScaling.EXP: 2 * PI / (2 ** (qubit_index + 1)),
     }
-    return conversion_dict[spectrum]
+    return conversion_dict[reupload_scaling]
 
 
-def basis_func(basis: str, x: Parameter) -> Parameter | sympy.Expr:
-    basis = basis.lower()
+def basis_func(basis: BasisSet, x: Parameter) -> Parameter | sympy.Expr:
     conversion_dict: dict[str, Parameter | sympy.Expr] = {
-        "fourier": x,
-        "chebyshev": 2 * sympy.acos(x),
+        BasisSet.FOURIER: x,
+        BasisSet.CHEBYSHEV: 2 * sympy.acos(x),
     }
     return conversion_dict[basis]
 
 
 def build_idx_fms(
-    basis: str,
+    basis: BasisSet,
     fm_pauli: Type[RY],
-    fm_strategy: str,
+    multivariate_strategy: MultivariateStrategy,
     n_features: int,
     n_qubits: int,
-    spectrum: str,
+    reupload_scaling: ReuploadScaling,
 ) -> list[KronBlock]:
     """Builds the index feature maps based on the given parameters.
 
     Args:
-        basis (str): Type of basis chosen for the feature map.
+        basis (BasisSet): Type of basis chosen for the feature map.
         fm_pauli (PrimitiveBlock type): The chosen Pauli rotation type.
-        fm_strategy (str): The feature map strategy to be used. Possible values are
-            'parallel' or 'serial'.
+        multivariate_strategy (MultivariateStrategy): The strategy used for encoding
+            the multivariate feature map.
         n_features (int): The number of features.
         n_qubits (int): The number of qubits.
-        spectrum (str): The chosen spectrum.
+        reupload_scaling (ReuploadScaling): The chosen scaling for the reupload.
 
     Returns:
         List[KronBlock]: The list of index feature maps.
     """
     idx_fms = []
     for i in range(n_features):
-        target_qubits = get_fm_qubits(fm_strategy, i, n_qubits, n_features)
+        target_qubits = get_fm_qubits(multivariate_strategy, i, n_qubits, n_features)
         param = FeatureParameter(f"x{i}")
         block = kron(
             *[
-                fm_pauli(qubit, generator_prefactor(spectrum, j) * basis_func(basis, param))
+                fm_pauli(qubit, generator_prefactor(reupload_scaling, j) * basis_func(basis, param))
                 for j, qubit in enumerate(target_qubits)
             ]
         )
@@ -70,12 +68,14 @@ def build_idx_fms(
     return idx_fms
 
 
-def get_fm_qubits(fm_strategy: str, i: int, n_qubits: int, n_features: int) -> Iterable:
+def get_fm_qubits(
+    multivariate_strategy: MultivariateStrategy, i: int, n_qubits: int, n_features: int
+) -> Iterable:
     """Returns the list of target qubits for the given feature map strategy and feature index.
 
     Args:
-        fm_strategy (str): The feature map strategy to be used. Possible values
-            are 'parallel' or 'serial'.
+        multivariate_strategy (MultivariateStrategy): The strategy used for encoding
+            the multivariate feature map.
         i (int): The feature index.
         n_qubits (int): The number of qubits.
         n_features (int): The number of features.
@@ -86,11 +86,11 @@ def get_fm_qubits(fm_strategy: str, i: int, n_qubits: int, n_features: int) -> I
     Raises:
         ValueError: If the feature map strategy is not implemented.
     """
-    if fm_strategy == "parallel":
+    if multivariate_strategy == MultivariateStrategy.PARALLEL:
         n_qubits_per_feature = int(n_qubits / n_features)
         target_qubits = range(i * n_qubits_per_feature, (i + 1) * n_qubits_per_feature)
-    elif fm_strategy == "serial":
+    elif multivariate_strategy == MultivariateStrategy.SERIES:
         target_qubits = range(0, n_qubits)
     else:
-        raise ValueError(f"Feature map strategy {fm_strategy} not implemented.")
+        raise ValueError(f"Multivariate strategy {multivariate_strategy} not implemented.")
     return target_qubits

qadence/engines/jax/differentiable_expectation.py

@@ -52,7 +52,7 @@ class DifferentiableExpectation:
             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()
+            k: tensor_to_jnp(v[0]) 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()