References

Ansaetze#

from qml_essentials.ansaetze import Ansaetze

Source code in qml_essentials/ansaetze.py

class Ansaetze:
    def get_available(parameterized_only=False):
        # list of parameterized ansaetze
        ansaetze = [
            Ansaetze.Circuit_1,
            Ansaetze.Circuit_2,
            Ansaetze.Circuit_3,
            Ansaetze.Circuit_4,
            Ansaetze.Circuit_5,
            Ansaetze.Circuit_6,
            Ansaetze.Circuit_7,
            Ansaetze.Circuit_8,
            Ansaetze.Circuit_9,
            Ansaetze.Circuit_10,
            Ansaetze.Circuit_13,
            Ansaetze.Circuit_14,
            Ansaetze.Circuit_15,
            Ansaetze.Circuit_16,
            Ansaetze.Circuit_17,
            Ansaetze.Circuit_18,
            Ansaetze.Circuit_19,
            Ansaetze.Circuit_20,
            Ansaetze.No_Entangling,
            Ansaetze.Strongly_Entangling,
            Ansaetze.Hardware_Efficient,
            Ansaetze.Permutation_Equivariant,
            Ansaetze.Matchgate,
            Ansaetze.XY_Brickwork,
        ]

        # extend by the non-parameterized ones
        if not parameterized_only:
            ansaetze += [
                Ansaetze.No_Ansatz,
                Ansaetze.GHZ,
            ]

        return ansaetze

    class No_Ansatz(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return ()

    class GHZ(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.H, wires=[0]),
                Block(
                    gate=Gates.CX,
                    topology=Topology.stairs,
                    reverse=False,
                    mirror=False,
                ),
            )

    class Circuit_1(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
            )

    class Circuit_2(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.CX,
                    topology=Topology.stairs,
                ),
            )

    class Circuit_3(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(gate=Gates.CRZ, topology=Topology.stairs),
            )

    class Circuit_4(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(gate=Gates.CRX, topology=Topology.stairs),
            )

    class Circuit_5(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(gate=Gates.CRZ, topology=Topology.all_to_all),
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
            )

    class Circuit_6(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(gate=Gates.CRX, topology=Topology.all_to_all),
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
            )

    class Circuit_7(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.CRZ,
                    topology=Topology.bricks,
                ),
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.CRZ,
                    topology=Topology.bricks,
                    offset=1,
                ),
            )

    class Circuit_8(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.CRX,
                    topology=Topology.bricks,
                ),
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.CRX,
                    topology=Topology.bricks,
                    offset=1,
                ),
            )

    class Circuit_9(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.H),
                Block(gate="CZ", topology=Topology.stairs),
                Block(gate=Gates.RX),
            )

    class Circuit_10(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RY),
                Block(gate="CZ", topology=Topology.stairs, offset=-1, wrap=True),
                Block(gate=Gates.RY),
            )

    class Circuit_13(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RY),
                Block(
                    gate=Gates.CRZ,
                    topology=Topology.stairs,
                    wrap=True,
                    reverse=True,
                    mirror=False,
                ),
                Block(gate=Gates.RY),
                Block(
                    gate=Gates.CRZ,
                    topology=Topology.stairs,
                    reverse=False,
                    mirror=False,
                    offset=lambda n: n - 1,
                    span=3,
                    wrap=True,
                ),
            )

    class Circuit_14(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RY),
                Block(
                    gate=Gates.CRX,
                    topology=Topology.stairs,
                    wrap=True,
                    reverse=True,
                    mirror=False,
                ),
                Block(gate=Gates.RY),
                Block(
                    gate=Gates.CRX,
                    topology=Topology.stairs,
                    reverse=False,
                    mirror=False,
                    offset=lambda n: n - 1,
                    span=3,
                    wrap=True,
                ),
            )

    class Circuit_15(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RY),
                Block(
                    gate=Gates.CX,
                    topology=Topology.stairs,
                    wrap=True,
                    reverse=True,
                    mirror=False,
                ),
                Block(gate=Gates.RY),
                Block(
                    gate=Gates.CX,
                    topology=Topology.stairs,
                    reverse=False,
                    mirror=False,
                    offset=lambda n: n - 1,
                    span=3,
                    wrap=True,
                ),
            )

    class Circuit_16(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.CRZ,
                    topology=Topology.bricks,
                ),
                Block(
                    gate=Gates.CRZ,
                    topology=Topology.bricks,
                    offset=1,
                ),
            )

    class Circuit_17(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.CRX,
                    topology=Topology.bricks,
                ),
                Block(
                    gate=Gates.CRX,
                    topology=Topology.bricks,
                    offset=1,
                ),
            )

    class Circuit_18(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.CRZ,
                    topology=Topology.stairs,
                    wrap=True,
                    mirror=False,
                ),
            )

    class Circuit_19(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX),
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.CRX,
                    topology=Topology.stairs,
                    wrap=True,
                    mirror=False,
                ),
            )

    class Circuit_20(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RY),
                Block(
                    gate=Gates.CX,
                    topology=Topology.stairs,
                    wrap=True,
                    reverse=True,
                    mirror=False,
                ),
                Block(gate=Gates.RY),
                Block(
                    gate=Gates.CX,
                    topology=Topology.stairs,
                    reverse=False,
                    offset=lambda n: n - 2,
                    span=1,
                    wrap=True,
                ),
            )

    class No_Entangling(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (Block(gate=Gates.Rot),)

    class Hardware_Efficient(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RY),
                Block(gate=Gates.RZ),
                Block(gate=Gates.RY),
                Block(
                    gate=Gates.CX,
                    topology=Topology.bricks,
                    mirror=False,
                ),
                Block(
                    gate=Gates.CX,
                    topology=Topology.bricks,
                    offset=-1,
                    modulo=True,
                    wrap=True,
                    mirror=False,
                ),
            )

    class Strongly_Entangling(DeclarativeCircuit):
        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.Rot),
                Block(
                    gate=Gates.CX,
                    topology=Topology.stairs,
                    wrap=True,
                    reverse=False,
                    mirror=False,
                ),
                Block(gate=Gates.Rot),
                Block(
                    gate=Gates.CX,
                    topology=Topology.stairs,
                    reverse=False,
                    span=lambda n: n // 2,
                    wrap=True,
                    mirror=False,
                ),
            )

    class Permutation_Equivariant(DeclarativeCircuit):
        r"""$S_n$ permutation-equivariant layer (Schatzki et al., arXiv:2210.09974).

        Shared-angle RX and RY on every qubit followed by a shared-angle RZZ on
        every qubit pair, realising $\exp(-i \frac{a}{2} \sum_k X_k)
        \exp(-i \frac{b}{2} \sum_k Y_k) \exp(-i \frac{c}{2} \sum_{j<k} Z_j Z_k)$
        for the rotation convention $R_P(\theta) = \exp(-i \frac{\theta}{2} P)$.
        The three parameters are tied (shared across all gates), so the layer
        width is 3 independent of the qubit count.

        Gradients assume JAX autodiff; a parameter-shift differentiator would need
        special handling for the shared parameters.
        """

        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RX, shared=True),
                Block(gate=Gates.RY, shared=True),
                Block(gate=Gates.RZZ, topology=Topology.all_pairs, shared=True),
            )

    class Matchgate(DeclarativeCircuit):
        r"""Matchgate LASA layer: RZ on every qubit + nearest-neighbour RXX.

        Generators $\{Z_k\} \cup \{X_k X_{k+1}\}$; the Lie closure is the
        matchgate algebra $\mathfrak{so}(2n)$ with $\dim = n(2n-1)$ (Kokcu et
        al., arXiv:2104.00728).  RXX is applied on the even nearest-neighbour
        bonds and then the odd bonds of the open chain, so the layer width is
        $n + (n-1)$.  Gradients assume JAX autodiff.
        """

        @classmethod
        def structure(cls):
            return (
                Block(gate=Gates.RZ),
                Block(
                    gate=Gates.RXX,
                    topology=Topology.bricks,
                    offset=0,
                    reverse=False,
                    mirror=False,
                ),
                Block(
                    gate=Gates.RXX,
                    topology=Topology.bricks,
                    offset=1,
                    reverse=False,
                    mirror=False,
                ),
            )

    class XY_Brickwork(DeclarativeCircuit):
        r"""Off-diagonal XY brickwork: nearest-neighbour RXX then RYY.

        Generators $\{X_k X_{k+1}, Y_k Y_{k+1}\}$; the Lie closure is the
        off-diagonal algebra $\mathfrak{so}(n) \oplus \mathfrak{so}(n)$ with no
        single-qubit $Z$, hence no deterministic $\mathfrak{g}$-purity floor.
        RXX on the even then odd bonds, followed by RYY on the even then odd
        bonds, so the layer width is $2(n-1)$.  Gradients assume JAX autodiff.
        """

        @classmethod
        def structure(cls):
            return (
                Block(
                    gate=Gates.RXX,
                    topology=Topology.bricks,
                    offset=0,
                    reverse=False,
                    mirror=False,
                ),
                Block(
                    gate=Gates.RXX,
                    topology=Topology.bricks,
                    offset=1,
                    reverse=False,
                    mirror=False,
                ),
                Block(
                    gate=Gates.RYY,
                    topology=Topology.bricks,
                    offset=0,
                    reverse=False,
                    mirror=False,
                ),
                Block(
                    gate=Gates.RYY,
                    topology=Topology.bricks,
                    offset=1,
                    reverse=False,
                    mirror=False,
                ),
            )

`Matchgate` #

Bases: DeclarativeCircuit

Matchgate LASA layer: RZ on every qubit + nearest-neighbour RXX.

Generators \(\{Z_k\} \cup \{X_k X_{k+1}\}\); the Lie closure is the matchgate algebra \(\mathfrak{so}(2n)\) with \(\dim = n(2n-1)\) (Kokcu et al., arXiv:2104.00728). RXX is applied on the even nearest-neighbour bonds and then the odd bonds of the open chain, so the layer width is \(n + (n-1)\). Gradients assume JAX autodiff.

Source code in qml_essentials/ansaetze.py

class Matchgate(DeclarativeCircuit):
    r"""Matchgate LASA layer: RZ on every qubit + nearest-neighbour RXX.

    Generators $\{Z_k\} \cup \{X_k X_{k+1}\}$; the Lie closure is the
    matchgate algebra $\mathfrak{so}(2n)$ with $\dim = n(2n-1)$ (Kokcu et
    al., arXiv:2104.00728).  RXX is applied on the even nearest-neighbour
    bonds and then the odd bonds of the open chain, so the layer width is
    $n + (n-1)$.  Gradients assume JAX autodiff.
    """

    @classmethod
    def structure(cls):
        return (
            Block(gate=Gates.RZ),
            Block(
                gate=Gates.RXX,
                topology=Topology.bricks,
                offset=0,
                reverse=False,
                mirror=False,
            ),
            Block(
                gate=Gates.RXX,
                topology=Topology.bricks,
                offset=1,
                reverse=False,
                mirror=False,
            ),
        )

`Permutation_Equivariant` #

Bases: DeclarativeCircuit

\(S_n\) permutation-equivariant layer (Schatzki et al., arXiv:2210.09974).

Shared-angle RX and RY on every qubit followed by a shared-angle RZZ on every qubit pair, realising \(\exp(-i \frac{a}{2} \sum_k X_k) \exp(-i \frac{b}{2} \sum_k Y_k) \exp(-i \frac{c}{2} \sum_{j<k} Z_j Z_k)\) for the rotation convention \(R_P(\theta) = \exp(-i \frac{\theta}{2} P)\). The three parameters are tied (shared across all gates), so the layer width is 3 independent of the qubit count.

Gradients assume JAX autodiff; a parameter-shift differentiator would need special handling for the shared parameters.

Source code in qml_essentials/ansaetze.py

class Permutation_Equivariant(DeclarativeCircuit):
    r"""$S_n$ permutation-equivariant layer (Schatzki et al., arXiv:2210.09974).

    Shared-angle RX and RY on every qubit followed by a shared-angle RZZ on
    every qubit pair, realising $\exp(-i \frac{a}{2} \sum_k X_k)
    \exp(-i \frac{b}{2} \sum_k Y_k) \exp(-i \frac{c}{2} \sum_{j<k} Z_j Z_k)$
    for the rotation convention $R_P(\theta) = \exp(-i \frac{\theta}{2} P)$.
    The three parameters are tied (shared across all gates), so the layer
    width is 3 independent of the qubit count.

    Gradients assume JAX autodiff; a parameter-shift differentiator would need
    special handling for the shared parameters.
    """

    @classmethod
    def structure(cls):
        return (
            Block(gate=Gates.RX, shared=True),
            Block(gate=Gates.RY, shared=True),
            Block(gate=Gates.RZZ, topology=Topology.all_pairs, shared=True),
        )

`XY_Brickwork` #

Bases: DeclarativeCircuit

Off-diagonal XY brickwork: nearest-neighbour RXX then RYY.

Generators \(\{X_k X_{k+1}, Y_k Y_{k+1}\}\); the Lie closure is the off-diagonal algebra \(\mathfrak{so}(n) \oplus \mathfrak{so}(n)\) with no single-qubit \(Z\), hence no deterministic \(\mathfrak{g}\)-purity floor. RXX on the even then odd bonds, followed by RYY on the even then odd bonds, so the layer width is \(2(n-1)\). Gradients assume JAX autodiff.

Source code in qml_essentials/ansaetze.py

class XY_Brickwork(DeclarativeCircuit):
    r"""Off-diagonal XY brickwork: nearest-neighbour RXX then RYY.

    Generators $\{X_k X_{k+1}, Y_k Y_{k+1}\}$; the Lie closure is the
    off-diagonal algebra $\mathfrak{so}(n) \oplus \mathfrak{so}(n)$ with no
    single-qubit $Z$, hence no deterministic $\mathfrak{g}$-purity floor.
    RXX on the even then odd bonds, followed by RYY on the even then odd
    bonds, so the layer width is $2(n-1)$.  Gradients assume JAX autodiff.
    """

    @classmethod
    def structure(cls):
        return (
            Block(
                gate=Gates.RXX,
                topology=Topology.bricks,
                offset=0,
                reverse=False,
                mirror=False,
            ),
            Block(
                gate=Gates.RXX,
                topology=Topology.bricks,
                offset=1,
                reverse=False,
                mirror=False,
            ),
            Block(
                gate=Gates.RYY,
                topology=Topology.bricks,
                offset=0,
                reverse=False,
                mirror=False,
            ),
            Block(
                gate=Gates.RYY,
                topology=Topology.bricks,
                offset=1,
                reverse=False,
                mirror=False,
            ),
        )

Circuit#

from qml_essentials.ansaetze import Circuit

Bases: ABC

Abstract base class for quantum circuit ansätze.

Source code in qml_essentials/ansaetze.py

class Circuit(ABC):
    """Abstract base class for quantum circuit ansätze."""

    def __init__(self) -> None:
        """Initialize the circuit."""
        pass

    @abstractmethod
    def n_params_per_layer(self, n_qubits: int) -> int:
        """
        Get the number of parameters per circuit layer.

        Args:
            n_qubits (int): Number of qubits in the circuit.

        Returns:
            int: Number of parameters required per layer.

        Raises:
            NotImplementedError: Must be implemented by subclasses.
        """
        raise NotImplementedError("n_params_per_layer method is not implemented")

    def n_pulse_params_per_layer(self, n_qubits: int) -> int:
        """
        Get the number of pulse parameters per circuit layer.

        Subclasses that do not use pulse-level simulation do not need to
        override this method.

        Args:
            n_qubits (int): Number of qubits in the circuit.

        Returns:
            int: Number of pulse parameters required per layer.

        Raises:
            NotImplementedError: If called but not overridden by subclass.
        """
        raise NotImplementedError("n_pulse_params_per_layer method is not implemented")

    @abstractmethod
    def get_control_indices(self, n_qubits: int) -> Optional[List[int]]:
        """
        Get indices for controlled rotation gates in one layer.

        Returns slice indices [start:stop:step] for extracting controlled
        gate parameters from a full parameter array for one layer.

        Args:
            n_qubits (int): Number of qubits in the circuit.

        Returns:
            Optional[List[int]]: List of three integers [start, stop, step]
                for slicing, or None if the circuit contains no controlled
                rotation gates.

        Raises:
            NotImplementedError: Must be implemented by subclasses.
        """
        raise NotImplementedError("get_control_indices method is not implemented")

    def get_control_angles(self, w: np.ndarray, n_qubits: int) -> Optional[np.ndarray]:
        """
        Extract angles for controlled rotation gates from parameter array.

        Args:
            w (np.ndarray): Parameter array for one layer.
            n_qubits (int): Number of qubits in the circuit.

        Returns:
            Optional[np.ndarray]: Array of controlled gate parameters,
                or empty array if circuit contains no controlled gates.
        """
        indices = self.get_control_indices(n_qubits)
        if indices is None:
            return np.array([])

        if len(indices) == 3 and None in indices:
            return w[indices[0] : indices[1] : indices[2]]
        else:
            return w.take(np.array(indices))

    def _build(self, w: np.ndarray, n_qubits: int, **kwargs: Any) -> Any:
        """
        Build one layer of the circuit using unitary or pulse-level parameters.

        Internal method that handles pulse parameter validation and context
        management before delegating to the build() method.

        Args:
            w (np.ndarray): Parameter array for the current layer.
            n_qubits (int): Number of qubits in the circuit.
            **kwargs: Additional keyword arguments:
                - pulse (bool): Whether to run the gates at pulse level.
                  Defaults to False.
                - pulse_params (np.ndarray): Pulse parameters if pulse=True.
                - noise_params (Dict): Noise parameters dictionary.

        Returns:
            Any: Result from the build() method.

        Raises:
            ValueError: If pulse_params length doesn't match expected count.
        """
        pulse = kwargs.get("pulse", False)

        if pulse and "pulse_params" in kwargs:
            pulse_params_per_layer = self.n_pulse_params_per_layer(n_qubits)

            if len(kwargs["pulse_params"]) != pulse_params_per_layer:
                raise ValueError(
                    f"Pulse params length {len(kwargs['pulse_params'])} "
                    f"does not match expected {pulse_params_per_layer} "
                    f"for {n_qubits} qubits"
                )

            with Gates.pulse_manager_context(kwargs["pulse_params"]):
                return self.build(w, n_qubits, **kwargs)
        else:
            return self.build(w, n_qubits, **kwargs)

    @abstractmethod
    def build(self, w: np.ndarray, n_qubits: int, **kwargs: Any) -> Any:
        """
        Build one layer of the quantum circuit.

        Args:
            w (np.ndarray): Parameter array for the current layer.
            n_qubits (int): Number of qubits in the circuit.
            **kwargs: Additional keyword arguments passed from _build.

        Returns:
            Any: Circuit construction result.

        Raises:
            NotImplementedError: Must be implemented by subclasses.
        """
        raise NotImplementedError("build method is not implemented")

    def __call__(self, *args: Any, **kwds: Any) -> Any:
        """Call the _build method with provided arguments."""
        self._build(*args, **kwds)

`call(*args, **kwds)` #

Call the _build method with provided arguments.

Source code in qml_essentials/ansaetze.py

def __call__(self, *args: Any, **kwds: Any) -> Any:
    """Call the _build method with provided arguments."""
    self._build(*args, **kwds)

`init()` #

Initialize the circuit.

Source code in qml_essentials/ansaetze.py

def __init__(self) -> None:
    """Initialize the circuit."""
    pass

`build(w, n_qubits, **kwargs)` `abstractmethod` #

Build one layer of the quantum circuit.

Parameters:

Name	Type	Description	Default
`w`	`ndarray`	Parameter array for the current layer.	required
`n_qubits`	`int`	Number of qubits in the circuit.	required
`**kwargs`	`Any`	Additional keyword arguments passed from _build.	`{}`

Returns:

Name	Type	Description
`Any`	`Any`	Circuit construction result.

Raises:

Type	Description
`NotImplementedError`	Must be implemented by subclasses.

Source code in qml_essentials/ansaetze.py

@abstractmethod
def build(self, w: np.ndarray, n_qubits: int, **kwargs: Any) -> Any:
    """
    Build one layer of the quantum circuit.

    Args:
        w (np.ndarray): Parameter array for the current layer.
        n_qubits (int): Number of qubits in the circuit.
        **kwargs: Additional keyword arguments passed from _build.

    Returns:
        Any: Circuit construction result.

    Raises:
        NotImplementedError: Must be implemented by subclasses.
    """
    raise NotImplementedError("build method is not implemented")

`get_control_angles(w, n_qubits)` #

Extract angles for controlled rotation gates from parameter array.

Parameters:

Name	Type	Description	Default
`w`	`ndarray`	Parameter array for one layer.	required
`n_qubits`	`int`	Number of qubits in the circuit.	required

Returns:

Type	Description
`Optional[ndarray]`	Optional[np.ndarray]: Array of controlled gate parameters, or empty array if circuit contains no controlled gates.

Source code in qml_essentials/ansaetze.py

def get_control_angles(self, w: np.ndarray, n_qubits: int) -> Optional[np.ndarray]:
    """
    Extract angles for controlled rotation gates from parameter array.

    Args:
        w (np.ndarray): Parameter array for one layer.
        n_qubits (int): Number of qubits in the circuit.

    Returns:
        Optional[np.ndarray]: Array of controlled gate parameters,
            or empty array if circuit contains no controlled gates.
    """
    indices = self.get_control_indices(n_qubits)
    if indices is None:
        return np.array([])

    if len(indices) == 3 and None in indices:
        return w[indices[0] : indices[1] : indices[2]]
    else:
        return w.take(np.array(indices))

`get_control_indices(n_qubits)` `abstractmethod` #

Get indices for controlled rotation gates in one layer.

Returns slice indices [start:stop:step] for extracting controlled gate parameters from a full parameter array for one layer.

Parameters:

Name	Type	Description	Default
`n_qubits`	`int`	Number of qubits in the circuit.	required

Returns:

Type	Description
`Optional[List[int]]`	Optional[List[int]]: List of three integers [start, stop, step] for slicing, or None if the circuit contains no controlled rotation gates.

Raises:

Type	Description
`NotImplementedError`	Must be implemented by subclasses.

Source code in qml_essentials/ansaetze.py

@abstractmethod
def get_control_indices(self, n_qubits: int) -> Optional[List[int]]:
    """
    Get indices for controlled rotation gates in one layer.

    Returns slice indices [start:stop:step] for extracting controlled
    gate parameters from a full parameter array for one layer.

    Args:
        n_qubits (int): Number of qubits in the circuit.

    Returns:
        Optional[List[int]]: List of three integers [start, stop, step]
            for slicing, or None if the circuit contains no controlled
            rotation gates.

    Raises:
        NotImplementedError: Must be implemented by subclasses.
    """
    raise NotImplementedError("get_control_indices method is not implemented")

`n_params_per_layer(n_qubits)` `abstractmethod` #

Get the number of parameters per circuit layer.

Parameters:

Name	Type	Description	Default
`n_qubits`	`int`	Number of qubits in the circuit.	required

Returns:

Name	Type	Description
`int`	`int`	Number of parameters required per layer.

Raises:

Type	Description
`NotImplementedError`	Must be implemented by subclasses.

Source code in qml_essentials/ansaetze.py

@abstractmethod
def n_params_per_layer(self, n_qubits: int) -> int:
    """
    Get the number of parameters per circuit layer.

    Args:
        n_qubits (int): Number of qubits in the circuit.

    Returns:
        int: Number of parameters required per layer.

    Raises:
        NotImplementedError: Must be implemented by subclasses.
    """
    raise NotImplementedError("n_params_per_layer method is not implemented")

`n_pulse_params_per_layer(n_qubits)` #

Get the number of pulse parameters per circuit layer.

Subclasses that do not use pulse-level simulation do not need to override this method.

Parameters:

Name	Type	Description	Default
`n_qubits`	`int`	Number of qubits in the circuit.	required

Returns:

Name	Type	Description
`int`	`int`	Number of pulse parameters required per layer.

Raises:

Type	Description
`NotImplementedError`	If called but not overridden by subclass.

Source code in qml_essentials/ansaetze.py

def n_pulse_params_per_layer(self, n_qubits: int) -> int:
    """
    Get the number of pulse parameters per circuit layer.

    Subclasses that do not use pulse-level simulation do not need to
    override this method.

    Args:
        n_qubits (int): Number of qubits in the circuit.

    Returns:
        int: Number of pulse parameters required per layer.

    Raises:
        NotImplementedError: If called but not overridden by subclass.
    """
    raise NotImplementedError("n_pulse_params_per_layer method is not implemented")

Declarative Circuit#

from qml_essentials.ansaetze import DeclarativeCircuit

Bases: Circuit

A circuit defined entirely by a sequence of Block descriptors.

Subclasses only need to set the class attribute structure — a tuple of

All of n_params_per_layer, n_pulse_params_per_layer, get_control_indices, and build are derived automatically.

Source code in qml_essentials/ansaetze.py

class DeclarativeCircuit(Circuit):
    """
    A circuit defined entirely by a sequence of Block descriptors.

    Subclasses only need to set the class attribute `structure` — a tuple of

    All of `n_params_per_layer`, `n_pulse_params_per_layer`,
    `get_control_indices`, and `build` are derived automatically.
    """

    @classmethod
    def structure(cls) -> Tuple[Any, ...]:
        """Override in subclass to return the structure tuple."""
        raise NotImplementedError

    @classmethod
    def n_params_per_layer(cls, n_qubits: int) -> int:
        return sum(block.n_params(n_qubits) for block in cls.structure())

    @classmethod
    def n_pulse_params_per_layer(cls, n_qubits: int) -> int:
        return sum(block.n_pulse_params(n_qubits) for block in cls.structure())

    @classmethod
    def get_control_indices(cls, n_qubits: int) -> Optional[List]:
        """
        Computes parameter indices for controlled rotation Gates.
        Scans the structure for Block with
        [start, stop, step] into the flat parameter vector, or None.
        """
        structure = cls.structure()
        total_params = sum(block.n_params(n_qubits) for block in structure)

        # Collect which parameter indices correspond to controlled rotations
        controlled_indices = []
        offset = 0
        for block in structure:
            n = block.n_params(n_qubits)
            if block.is_controlled_rotation:
                controlled_indices.extend(range(offset, offset + n))
            offset += n

        # FIXME: this last part should be reworked

        if not controlled_indices:
            return None

        # Check if indices form a contiguous tail (the common case)
        # This preserves backwards compatibility with the [start, None, None] format
        if controlled_indices == list(
            range(total_params - len(controlled_indices), total_params)
        ):
            return [-len(controlled_indices), None, None]

        # Fallback: return raw indices (future-proof)
        return controlled_indices

    @classmethod
    def build(cls, w: np.ndarray, n_qubits: int, **kwargs: Any) -> None:
        structure = cls.structure()
        w_idx = 0
        for block in structure:
            w_idx = block.apply(n_qubits, w, w_idx, **kwargs)
            Gates.Barrier(wires=list(range(n_qubits)), **kwargs)

`get_control_indices(n_qubits)` `classmethod` #

Computes parameter indices for controlled rotation Gates. Scans the structure for Block with [start, stop, step] into the flat parameter vector, or None.

Source code in qml_essentials/ansaetze.py

@classmethod
def get_control_indices(cls, n_qubits: int) -> Optional[List]:
    """
    Computes parameter indices for controlled rotation Gates.
    Scans the structure for Block with
    [start, stop, step] into the flat parameter vector, or None.
    """
    structure = cls.structure()
    total_params = sum(block.n_params(n_qubits) for block in structure)

    # Collect which parameter indices correspond to controlled rotations
    controlled_indices = []
    offset = 0
    for block in structure:
        n = block.n_params(n_qubits)
        if block.is_controlled_rotation:
            controlled_indices.extend(range(offset, offset + n))
        offset += n

    # FIXME: this last part should be reworked

    if not controlled_indices:
        return None

    # Check if indices form a contiguous tail (the common case)
    # This preserves backwards compatibility with the [start, None, None] format
    if controlled_indices == list(
        range(total_params - len(controlled_indices), total_params)
    ):
        return [-len(controlled_indices), None, None]

    # Fallback: return raw indices (future-proof)
    return controlled_indices

`structure()` `classmethod` #

Override in subclass to return the structure tuple.

Source code in qml_essentials/ansaetze.py

@classmethod
def structure(cls) -> Tuple[Any, ...]:
    """Override in subclass to return the structure tuple."""
    raise NotImplementedError

Block#

from qml_essentials.ansaetze import Block

Source code in qml_essentials/ansaetze.py

class Block:
    def __init__(
        self,
        gate: str,
        topology: Any = None,
        shared: bool = False,
        wires: Optional[List[int]] = None,
        **kwargs,
    ):
        """
        Initialize a Block object; the atoms of Ansatzes.

        Args:
            gate (str): Name of the Gate class to use.
            topology (Any, optional): Topology of the gate for entangling gates.
                Defaults to None.
            shared (bool, optional): Tie all gates of the block to a single
                parameter (per-gate width), instead of one parameter per gate.
                Defaults to False.
            wires (Optional[List[int]], optional): Fixed wires for a
                non-entangling block. If None, the block spans all qubits.
                Defaults to None.
            kwargs (Any): Additional keyword arguments passed to the topology function.
        """
        if isinstance(gate, str):
            self.gate = getattr(Gates, gate)
        else:
            self.gate = gate

        if self.is_entangling:
            assert topology is not None, (
                "Topology must be specified for entangling gates"
            )

        self.topology = topology
        self.shared = shared
        self.wires = wires
        self.kwargs = kwargs

    def __repr__(self):
        if self.topology is None:
            return f"{self.__class__.__name__}({self.gate.__name__})"
        else:
            return (
                f"{self.__class__.__name__}"
                f"({self.topology.__name__}[{self.gate.__name__}])"
            )

    @property
    def is_entangling(self):
        return Gates.is_entangling(self.gate)

    @property
    def is_rotational(self):
        return Gates.is_rotational(self.gate)

    @property
    def is_controlled_rotation(self):
        return Gates.is_controlled(self.gate) and self.is_rotational

    def enough_qubits(self, n_qubits):
        if self.is_entangling:
            # NOTE This must be adjusted if default values
            # in Topology change
            span = self.kwargs.get("span", 1)
            if callable(span):
                span = span(n_qubits)

            return (n_qubits >= 2) and (n_qubits > span)

        return n_qubits >= 1

    def n_params(self, n_qubits: int) -> int:
        assert n_qubits > 0, "Number of qubits must be positive"

        if not self.is_rotational:
            return 0

        per_gate = 3 if self.gate.__name__ == "Rot" else 1

        if self.is_entangling:
            if not self.enough_qubits(n_qubits):
                warnings.warn(
                    f"Skipping {self.topology.__name__} with n_qubits={n_qubits} "
                    f"as there are not enough qubits"
                    f"for this topology."
                )
                return 0
            n_gates = len(self.topology(n_qubits=n_qubits, **self.kwargs))
        else:
            n_gates = len(self.wires) if self.wires is not None else n_qubits

        if n_gates == 0:  # an empty block consumes no parameters, shared or not
            return 0

        return per_gate if self.shared else per_gate * n_gates

    def n_pulse_params(self, n_qubits: int) -> int:
        assert n_qubits > 0, "Number of qubits must be positive"

        n_pulse_params = PulseInformation.num_params(self.gate)
        if self.is_entangling:
            if not self.enough_qubits(n_qubits):
                warnings.warn(
                    f"Skipping {self.topology.__name__} with n_qubits={n_qubits} "
                    f"as there are not enough qubits"
                    f"for this topology."
                )
                return 0
            else:
                return n_pulse_params * len(
                    self.topology(n_qubits=n_qubits, **self.kwargs)
                )
        n_gates = len(self.wires) if self.wires is not None else n_qubits
        return n_pulse_params * n_gates

    def apply(
        self, n_qubits: int, w: np.ndarray = None, w_idx: int = None, **kwargs
    ) -> int:
        """
        Applies the block to the given circuit.

        Args:
            n_qubits (int): Number of qubits, the block is applied to.
            w (np.ndarray, optional): Weights to use for rotational gates.
                Defaults to None.
            w_idx (int, optional): Index of weights to use for rotational gates.
                Defaults to None.
            **kwargs (Any): Keyword arguments passed to the gate.

        Returns:
            int: The new index of weights after applying the block.
        """
        assert n_qubits > 0, "Number of qubits must be positive"

        if self.is_entangling:
            iterator = self.topology(n_qubits=n_qubits, **self.kwargs)
        else:
            iterator = self.wires if self.wires is not None else range(n_qubits)

        per_gate = 3 if self.gate.__name__ == "Rot" else 1
        base = w_idx  # start index, reused for every gate when shared
        applied = False

        for wires in iterator:
            if self.is_entangling and not self.enough_qubits(n_qubits):
                warnings.warn(
                    f"Skipping {self.topology.__name__} with n_qubits={n_qubits} "
                    f"as there are not enough qubits"
                    f"for this topology."
                )
                continue

            if self.is_rotational:
                assert w is not None, "w must be provided for rotational gates"
                assert w_idx is not None, "w_idx must be provided for rotational gates"

                i = base if self.shared else w_idx
                if per_gate == 3:
                    self.gate(w[i], w[i + 1], w[i + 2], wires=wires, **kwargs)
                else:
                    self.gate(w[i], wires=wires, **kwargs)
                if not self.shared:
                    w_idx += per_gate
                applied = True
            else:
                self.gate(wires=wires, **kwargs)

        if self.is_rotational and self.shared and applied:
            w_idx = base + per_gate
        return w_idx

`init(gate, topology=None, shared=False, wires=None, **kwargs)` #

Initialize a Block object; the atoms of Ansatzes.

Parameters:

Name	Type	Description	Default
`gate`	`str`	Name of the Gate class to use.	required
`topology`	`Any`	Topology of the gate for entangling gates. Defaults to None.	`None`
`shared`	`bool`	Tie all gates of the block to a single parameter (per-gate width), instead of one parameter per gate. Defaults to False.	`False`
`wires`	`Optional[List[int]]`	Fixed wires for a non-entangling block. If None, the block spans all qubits. Defaults to None.	`None`
`kwargs`	`Any`	Additional keyword arguments passed to the topology function.	`{}`

Source code in qml_essentials/ansaetze.py

def __init__(
    self,
    gate: str,
    topology: Any = None,
    shared: bool = False,
    wires: Optional[List[int]] = None,
    **kwargs,
):
    """
    Initialize a Block object; the atoms of Ansatzes.

    Args:
        gate (str): Name of the Gate class to use.
        topology (Any, optional): Topology of the gate for entangling gates.
            Defaults to None.
        shared (bool, optional): Tie all gates of the block to a single
            parameter (per-gate width), instead of one parameter per gate.
            Defaults to False.
        wires (Optional[List[int]], optional): Fixed wires for a
            non-entangling block. If None, the block spans all qubits.
            Defaults to None.
        kwargs (Any): Additional keyword arguments passed to the topology function.
    """
    if isinstance(gate, str):
        self.gate = getattr(Gates, gate)
    else:
        self.gate = gate

    if self.is_entangling:
        assert topology is not None, (
            "Topology must be specified for entangling gates"
        )

    self.topology = topology
    self.shared = shared
    self.wires = wires
    self.kwargs = kwargs

`apply(n_qubits, w=None, w_idx=None, **kwargs)` #

Applies the block to the given circuit.

Parameters:

Name	Type	Description	Default
`n_qubits`	`int`	Number of qubits, the block is applied to.	required
`w`	`ndarray`	Weights to use for rotational gates. Defaults to None.	`None`
`w_idx`	`int`	Index of weights to use for rotational gates. Defaults to None.	`None`
`**kwargs`	`Any`	Keyword arguments passed to the gate.	`{}`

Returns:

Name	Type	Description
`int`	`int`	The new index of weights after applying the block.

Source code in qml_essentials/ansaetze.py

def apply(
    self, n_qubits: int, w: np.ndarray = None, w_idx: int = None, **kwargs
) -> int:
    """
    Applies the block to the given circuit.

    Args:
        n_qubits (int): Number of qubits, the block is applied to.
        w (np.ndarray, optional): Weights to use for rotational gates.
            Defaults to None.
        w_idx (int, optional): Index of weights to use for rotational gates.
            Defaults to None.
        **kwargs (Any): Keyword arguments passed to the gate.

    Returns:
        int: The new index of weights after applying the block.
    """
    assert n_qubits > 0, "Number of qubits must be positive"

    if self.is_entangling:
        iterator = self.topology(n_qubits=n_qubits, **self.kwargs)
    else:
        iterator = self.wires if self.wires is not None else range(n_qubits)

    per_gate = 3 if self.gate.__name__ == "Rot" else 1
    base = w_idx  # start index, reused for every gate when shared
    applied = False

    for wires in iterator:
        if self.is_entangling and not self.enough_qubits(n_qubits):
            warnings.warn(
                f"Skipping {self.topology.__name__} with n_qubits={n_qubits} "
                f"as there are not enough qubits"
                f"for this topology."
            )
            continue

        if self.is_rotational:
            assert w is not None, "w must be provided for rotational gates"
            assert w_idx is not None, "w_idx must be provided for rotational gates"

            i = base if self.shared else w_idx
            if per_gate == 3:
                self.gate(w[i], w[i + 1], w[i + 2], wires=wires, **kwargs)
            else:
                self.gate(w[i], wires=wires, **kwargs)
            if not self.shared:
                w_idx += per_gate
            applied = True
        else:
            self.gate(wires=wires, **kwargs)

    if self.is_rotational and self.shared and applied:
        w_idx = base + per_gate
    return w_idx

Encoding#

from qml_essentials.ansaetze import Encoding

Source code in qml_essentials/ansaetze.py

class Encoding:
    def __init__(
        self, strategy: str, gates: Union[str, Callable, List[Union[str, Callable]]]
    ):
        """
        Initializes an Encoding object.

        Implementations closely follow https://doi.org/10.22331/q-2023-12-20-1210

        Parameters
        ----------
        strategy : str
            The encoding strategy to use. Available options:
            ['hamming', 'binary', 'ternary']
        gates : Union[str, Callable, List[Union[str, Callable]]]
            The gates to use for encoding. Can be a string, a callable or a list
            of strings or callables.

        Returns
        -------
        None

        Raises
        -------
        ValueError
            If the encoding strategy is not implemented.
        ValueError
            If there is an error parsing the Gates.
        """
        if strategy not in ["hamming", "binary", "ternary", "golomb"]:
            raise ValueError(
                f"Encoding strategy {strategy} not implemented. "
                "Available options: ['hamming', 'binary', 'ternary', 'golomb']"
            )
        self._strategy = strategy
        strategy_fn = getattr(self, strategy)

        log.debug(f"Using encoding strategy: '{strategy_fn.__name__}'")

        if self._strategy == "golomb":
            self._gates = []
            self.callable = [strategy_fn(None)]
        else:
            try:
                self._gates = Gates.parse_gates(gates, Gates)
            except ValueError as e:
                raise ValueError(f"Error parsing encodings: {e}")

            self.callable = [strategy_fn(g) for g in self._gates]

    def __len__(self):
        return len(self.callable)

    def __getitem__(self, idx):
        return self.callable[idx]

    def get_n_freqs(self, data_reupload):
        """
        Number of reachable frequencies (positive + negative + DC) for the
        encoding strategy, given the ``(n_layers, n_qubits)`` data-reupload mask.
        """
        return int(self.get_spectrum(data_reupload).size)

    def get_spectrum(self, data_reupload):
        """
        Reachable Fourier frequency comb for the encoding strategy.

        Computed exactly from the ``(n_layers, n_qubits)`` data-reupload mask as
        the Minkowski sum of the per-gate generator frequencies:

        - hamming: every encoding gate contributes +/-1, so the comb is
          ``{-k, ..., k}`` with ``k`` the total number of encoding gates.
        - binary / ternary: qubit ``q`` is scaled by ``base**q`` (base 2 / 3),
          applied once per active layer, so the comb is the Minkowski sum over
          qubits of ``{k * base**q : |k| <= count_q}`` with ``count_q`` the
          number of layers that re-upload on qubit ``q``.
        - golomb: a single multi-qubit diagonal gate per *active layer* (see
          ``Model._iec``), each spanning ``[-max_mark, max_mark]``; ``k`` active
          layers give ``{-k*max_mark, ..., k*max_mark}``. The contiguous range
          is returned (not the sparse mark-difference set) because the FFT in
          ``Coefficients._fourier_transform`` samples at ``model.degree``
          resolution and must cover the max frequency; residual sparse gaps
          carry ~0 coefficients.

        See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

        Parameters
        ----------
        data_reupload : np.ndarray
            Boolean mask of shape ``(n_layers, n_qubits)`` (or ``(n_qubits,)``
            for a single layer) marking where the encoding re-uploads.

        Returns
        -------
        np.ndarray
            The sorted reachable spectrum of the encoding strategy.
        """
        mask = np.asarray(data_reupload, dtype=bool)
        if mask.ndim == 1:  # (n_qubits,) -> treat as a single layer
            mask = mask[None, :]

        if self._strategy not in ("hamming", "binary", "ternary", "golomb"):
            raise NotImplementedError
        if self._strategy == "golomb":
            n_qubits = getattr(self, "_n_qubits", None)
            if n_qubits is None:
                raise ValueError("Golomb encoding requires n_qubits to be set")
            apps = int(np.count_nonzero(mask.any(axis=1)))  # one gate per active layer
            limit = apps * max(golomb_ruler(2**n_qubits))
            return np.arange(-limit, limit + 1)

        base = {"hamming": 1, "binary": 2, "ternary": 3}[self._strategy]
        counts = mask.sum(axis=0)  # per-qubit re-upload count (index == wire)
        reach = {0}
        for q, c in enumerate(counts):
            scale = base**q
            reach = {
                a + k
                for a in reach
                for k in range(-int(c) * scale, int(c) * scale + 1, scale)
            }
        return np.array(sorted(reach))

    def get_weights(self, n_qubits):
        """
        Per-qubit weight vector w for the separable weighted encodings.

        The encoding loads the scaled input phi_q = w_q * x on qubit q, so the
        returned weights match the per-qubit scaling of the strategy callables
        (see :meth:`binary` and :meth:`ternary`).

        Parameters
        ----------
        n_qubits : int
            The number of qubits carrying the encoding.

        Returns
        -------
        np.ndarray
            The weight vector of shape ``(n_qubits,)``.

        Raises
        ------
        ValueError
            If the strategy is non-separable (golomb) and has no per-qubit weights.
        """
        if self._strategy == "hamming":
            return np.ones(n_qubits)
        elif self._strategy == "binary":
            return 2.0 ** np.arange(n_qubits)
        elif self._strategy == "ternary":
            return 3.0 ** np.arange(n_qubits)
        elif self._strategy == "golomb":
            raise ValueError(
                "Golomb encoding is non-separable and has no per-qubit weights."
            )
        else:
            raise NotImplementedError

    def hamming(self, enc):
        """
        Hamming encoding strategy.

        Returns an encoding function that uses the Hamming encoding strategy
        which uses 2 * omegas + 1 frequencies for the encoding.
        See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

        Parameters
        ----------
        enc : Callable
            The encoding function to be wrapped.

        Returns
        -------
        Callable
            The wrapped encoding function.
        """
        return enc

    def binary(self, enc):
        """
        Binary encoding strategy.

        Returns an encoding function that scales the input by a factor of 2^wires.

        Binary encoding uses 2^(omegas + 1) - 1 frequencies for the encoding.
        See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

        Parameters
        ----------
        enc : Callable
            The encoding function to be wrapped.

        Returns
        -------
        Callable
            The wrapped encoding function.
        """

        def _enc(inputs, wires, **kwargs):
            return enc(inputs * (2**wires), wires, **kwargs)

        return _enc

    def ternary(self, enc):
        """
        Ternary encoding strategy.

        Returns an encoding function that scales the input by a factor of 3^wires.

        Ternary encoding uses 3^omegas frequencies for the encoding.
        See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

        Parameters
        ----------
        enc : Callable
            The encoding function to be wrapped.

        Returns
        -------
        Callable
            The wrapped encoding function.
        """

        def _enc(inputs, wires, **kwargs):
            return enc(inputs * (3**wires), wires, **kwargs)

        return _enc

    @property
    def is_golomb(self):
        """Whether this encoding uses the Golomb (multi-qubit diagonal) strategy."""
        return self._strategy == "golomb"

    def golomb(self, enc):
        """Golomb encoding strategy.

        Returns a callable that applies a multi-qubit diagonal unitary
        ``S(x) = exp(-i H x)`` where ``H = diag(golomb_marks)`` to all
        qubits simultaneously.  This produces the largest possible
        ``|Ω| = d(d-1)+1`` for any *d*-dimensional Hamiltonian, with
        ``|R(k)| = 1`` for all nonzero frequencies *k*.

        Unlike the other strategies, Golomb encoding does *not* wrap a
        per-qubit gate.  Instead, the model's ``_iec`` method detects
        ``is_golomb`` and applies a single ``GolombEncoding`` gate on
        all qubits.

        See Peters et al., arXiv:2209.05523, Sec. 3.1 and Appendix C.4.

        Parameters
        ----------
        enc : Callable or None
            Ignored (Golomb encoding uses its own multi-qubit gate).

        Returns
        -------
        Callable
            A callable with the same signature as per-qubit encoding
            functions but that applies :func:`GolombEncoding`.
        """

        def _enc(inputs, wires, **kwargs):
            # `wires` here is a list of all qubit indices, set by _iec
            GolombEncoding(w=inputs, wires=wires, **kwargs)

        return _enc

`is_golomb` `property` #

Whether this encoding uses the Golomb (multi-qubit diagonal) strategy.

`init(strategy, gates)` #

Initializes an Encoding object.

Implementations closely follow https://doi.org/10.22331/q-2023-12-20-1210

Parameters#

strategy : str The encoding strategy to use. Available options: ['hamming', 'binary', 'ternary'] gates : Union[str, Callable, List[Union[str, Callable]]] The gates to use for encoding. Can be a string, a callable or a list of strings or callables.

Returns#

None

Raises#

ValueError If the encoding strategy is not implemented. ValueError If there is an error parsing the Gates.

Source code in qml_essentials/ansaetze.py

def __init__(
    self, strategy: str, gates: Union[str, Callable, List[Union[str, Callable]]]
):
    """
    Initializes an Encoding object.

    Implementations closely follow https://doi.org/10.22331/q-2023-12-20-1210

    Parameters
    ----------
    strategy : str
        The encoding strategy to use. Available options:
        ['hamming', 'binary', 'ternary']
    gates : Union[str, Callable, List[Union[str, Callable]]]
        The gates to use for encoding. Can be a string, a callable or a list
        of strings or callables.

    Returns
    -------
    None

    Raises
    -------
    ValueError
        If the encoding strategy is not implemented.
    ValueError
        If there is an error parsing the Gates.
    """
    if strategy not in ["hamming", "binary", "ternary", "golomb"]:
        raise ValueError(
            f"Encoding strategy {strategy} not implemented. "
            "Available options: ['hamming', 'binary', 'ternary', 'golomb']"
        )
    self._strategy = strategy
    strategy_fn = getattr(self, strategy)

    log.debug(f"Using encoding strategy: '{strategy_fn.__name__}'")

    if self._strategy == "golomb":
        self._gates = []
        self.callable = [strategy_fn(None)]
    else:
        try:
            self._gates = Gates.parse_gates(gates, Gates)
        except ValueError as e:
            raise ValueError(f"Error parsing encodings: {e}")

        self.callable = [strategy_fn(g) for g in self._gates]

`binary(enc)` #

Binary encoding strategy.

Returns an encoding function that scales the input by a factor of 2^wires.

Binary encoding uses 2^(omegas + 1) - 1 frequencies for the encoding. See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

Parameters#

enc : Callable The encoding function to be wrapped.

Returns#

Callable The wrapped encoding function.

Source code in qml_essentials/ansaetze.py

def binary(self, enc):
    """
    Binary encoding strategy.

    Returns an encoding function that scales the input by a factor of 2^wires.

    Binary encoding uses 2^(omegas + 1) - 1 frequencies for the encoding.
    See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

    Parameters
    ----------
    enc : Callable
        The encoding function to be wrapped.

    Returns
    -------
    Callable
        The wrapped encoding function.
    """

    def _enc(inputs, wires, **kwargs):
        return enc(inputs * (2**wires), wires, **kwargs)

    return _enc

`get_n_freqs(data_reupload)` #

Number of reachable frequencies (positive + negative + DC) for the encoding strategy, given the (n_layers, n_qubits) data-reupload mask.

Source code in qml_essentials/ansaetze.py

def get_n_freqs(self, data_reupload):
    """
    Number of reachable frequencies (positive + negative + DC) for the
    encoding strategy, given the ``(n_layers, n_qubits)`` data-reupload mask.
    """
    return int(self.get_spectrum(data_reupload).size)

`get_spectrum(data_reupload)` #

Reachable Fourier frequency comb for the encoding strategy.

Computed exactly from the (n_layers, n_qubits) data-reupload mask as the Minkowski sum of the per-gate generator frequencies:

hamming: every encoding gate contributes +/-1, so the comb is {-k, ..., k} with k the total number of encoding gates.
binary / ternary: qubit q is scaled by base**q (base 2 / 3), applied once per active layer, so the comb is the Minkowski sum over qubits of {k * base**q : |k| <= count_q} with count_q the number of layers that re-upload on qubit q.
golomb: a single multi-qubit diagonal gate per active layer (see Model._iec), each spanning [-max_mark, max_mark]; k active layers give {-k*max_mark, ..., k*max_mark}. The contiguous range is returned (not the sparse mark-difference set) because the FFT in Coefficients._fourier_transform samples at model.degree resolution and must cover the max frequency; residual sparse gaps carry ~0 coefficients.

See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

Parameters#

data_reupload : np.ndarray Boolean mask of shape (n_layers, n_qubits) (or (n_qubits,) for a single layer) marking where the encoding re-uploads.

Returns#

np.ndarray The sorted reachable spectrum of the encoding strategy.

Source code in qml_essentials/ansaetze.py

def get_spectrum(self, data_reupload):
    """
    Reachable Fourier frequency comb for the encoding strategy.

    Computed exactly from the ``(n_layers, n_qubits)`` data-reupload mask as
    the Minkowski sum of the per-gate generator frequencies:

    - hamming: every encoding gate contributes +/-1, so the comb is
      ``{-k, ..., k}`` with ``k`` the total number of encoding gates.
    - binary / ternary: qubit ``q`` is scaled by ``base**q`` (base 2 / 3),
      applied once per active layer, so the comb is the Minkowski sum over
      qubits of ``{k * base**q : |k| <= count_q}`` with ``count_q`` the
      number of layers that re-upload on qubit ``q``.
    - golomb: a single multi-qubit diagonal gate per *active layer* (see
      ``Model._iec``), each spanning ``[-max_mark, max_mark]``; ``k`` active
      layers give ``{-k*max_mark, ..., k*max_mark}``. The contiguous range
      is returned (not the sparse mark-difference set) because the FFT in
      ``Coefficients._fourier_transform`` samples at ``model.degree``
      resolution and must cover the max frequency; residual sparse gaps
      carry ~0 coefficients.

    See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

    Parameters
    ----------
    data_reupload : np.ndarray
        Boolean mask of shape ``(n_layers, n_qubits)`` (or ``(n_qubits,)``
        for a single layer) marking where the encoding re-uploads.

    Returns
    -------
    np.ndarray
        The sorted reachable spectrum of the encoding strategy.
    """
    mask = np.asarray(data_reupload, dtype=bool)
    if mask.ndim == 1:  # (n_qubits,) -> treat as a single layer
        mask = mask[None, :]

    if self._strategy not in ("hamming", "binary", "ternary", "golomb"):
        raise NotImplementedError
    if self._strategy == "golomb":
        n_qubits = getattr(self, "_n_qubits", None)
        if n_qubits is None:
            raise ValueError("Golomb encoding requires n_qubits to be set")
        apps = int(np.count_nonzero(mask.any(axis=1)))  # one gate per active layer
        limit = apps * max(golomb_ruler(2**n_qubits))
        return np.arange(-limit, limit + 1)

    base = {"hamming": 1, "binary": 2, "ternary": 3}[self._strategy]
    counts = mask.sum(axis=0)  # per-qubit re-upload count (index == wire)
    reach = {0}
    for q, c in enumerate(counts):
        scale = base**q
        reach = {
            a + k
            for a in reach
            for k in range(-int(c) * scale, int(c) * scale + 1, scale)
        }
    return np.array(sorted(reach))

`get_weights(n_qubits)` #

Per-qubit weight vector w for the separable weighted encodings.

The encoding loads the scaled input phi_q = w_q * x on qubit q, so the returned weights match the per-qubit scaling of the strategy callables (see :meth:binary and :meth:ternary).

Parameters#

n_qubits : int The number of qubits carrying the encoding.

Returns#

np.ndarray The weight vector of shape (n_qubits,).

Raises#

ValueError If the strategy is non-separable (golomb) and has no per-qubit weights.

Source code in qml_essentials/ansaetze.py

def get_weights(self, n_qubits):
    """
    Per-qubit weight vector w for the separable weighted encodings.

    The encoding loads the scaled input phi_q = w_q * x on qubit q, so the
    returned weights match the per-qubit scaling of the strategy callables
    (see :meth:`binary` and :meth:`ternary`).

    Parameters
    ----------
    n_qubits : int
        The number of qubits carrying the encoding.

    Returns
    -------
    np.ndarray
        The weight vector of shape ``(n_qubits,)``.

    Raises
    ------
    ValueError
        If the strategy is non-separable (golomb) and has no per-qubit weights.
    """
    if self._strategy == "hamming":
        return np.ones(n_qubits)
    elif self._strategy == "binary":
        return 2.0 ** np.arange(n_qubits)
    elif self._strategy == "ternary":
        return 3.0 ** np.arange(n_qubits)
    elif self._strategy == "golomb":
        raise ValueError(
            "Golomb encoding is non-separable and has no per-qubit weights."
        )
    else:
        raise NotImplementedError

`golomb(enc)` #

Golomb encoding strategy.

Returns a callable that applies a multi-qubit diagonal unitary S(x) = exp(-i H x) where H = diag(golomb_marks) to all qubits simultaneously. This produces the largest possible |Ω| = d(d-1)+1 for any d-dimensional Hamiltonian, with |R(k)| = 1 for all nonzero frequencies k.

Unlike the other strategies, Golomb encoding does not wrap a per-qubit gate. Instead, the model's _iec method detects is_golomb and applies a single GolombEncoding gate on all qubits.

See Peters et al., arXiv:2209.05523, Sec. 3.1 and Appendix C.4.

Parameters#

enc : Callable or None Ignored (Golomb encoding uses its own multi-qubit gate).

Returns#

Callable A callable with the same signature as per-qubit encoding functions but that applies :func:GolombEncoding.

Source code in qml_essentials/ansaetze.py

def golomb(self, enc):
    """Golomb encoding strategy.

    Returns a callable that applies a multi-qubit diagonal unitary
    ``S(x) = exp(-i H x)`` where ``H = diag(golomb_marks)`` to all
    qubits simultaneously.  This produces the largest possible
    ``|Ω| = d(d-1)+1`` for any *d*-dimensional Hamiltonian, with
    ``|R(k)| = 1`` for all nonzero frequencies *k*.

    Unlike the other strategies, Golomb encoding does *not* wrap a
    per-qubit gate.  Instead, the model's ``_iec`` method detects
    ``is_golomb`` and applies a single ``GolombEncoding`` gate on
    all qubits.

    See Peters et al., arXiv:2209.05523, Sec. 3.1 and Appendix C.4.

    Parameters
    ----------
    enc : Callable or None
        Ignored (Golomb encoding uses its own multi-qubit gate).

    Returns
    -------
    Callable
        A callable with the same signature as per-qubit encoding
        functions but that applies :func:`GolombEncoding`.
    """

    def _enc(inputs, wires, **kwargs):
        # `wires` here is a list of all qubit indices, set by _iec
        GolombEncoding(w=inputs, wires=wires, **kwargs)

    return _enc

`hamming(enc)` #

Hamming encoding strategy.

Returns an encoding function that uses the Hamming encoding strategy which uses 2 * omegas + 1 frequencies for the encoding. See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

Parameters#

enc : Callable The encoding function to be wrapped.

Returns#

Callable The wrapped encoding function.

Source code in qml_essentials/ansaetze.py

def hamming(self, enc):
    """
    Hamming encoding strategy.

    Returns an encoding function that uses the Hamming encoding strategy
    which uses 2 * omegas + 1 frequencies for the encoding.
    See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

    Parameters
    ----------
    enc : Callable
        The encoding function to be wrapped.

    Returns
    -------
    Callable
        The wrapped encoding function.
    """
    return enc

`ternary(enc)` #

Ternary encoding strategy.

Returns an encoding function that scales the input by a factor of 3^wires.

Ternary encoding uses 3^omegas frequencies for the encoding. See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

Parameters#

enc : Callable The encoding function to be wrapped.

Returns#

Callable The wrapped encoding function.

Source code in qml_essentials/ansaetze.py

def ternary(self, enc):
    """
    Ternary encoding strategy.

    Returns an encoding function that scales the input by a factor of 3^wires.

    Ternary encoding uses 3^omegas frequencies for the encoding.
    See https://doi.org/10.22331/q-2023-12-20-1210 for more details.

    Parameters
    ----------
    enc : Callable
        The encoding function to be wrapped.

    Returns
    -------
    Callable
        The wrapped encoding function.
    """

    def _enc(inputs, wires, **kwargs):
        return enc(inputs * (3**wires), wires, **kwargs)

    return _enc

Model#

from qml_essentials.model import Model

A quantum circuit model.

Source code in qml_essentials/model.py

class Model:
    """
    A quantum circuit model.
    """

    def __init__(
        self,
        n_qubits: int,
        n_layers: int,
        circuit_type: Union[str, type[Circuit]] = "No_Ansatz",
        data_reupload: Union[
            bool, List[List[bool]], List[List[List[bool]]], np.ndarray
        ] = True,
        state_preparation: Union[
            str, Callable, List[Union[str, Callable]], None
        ] = None,
        encoding: Union[Encoding, str, Callable, List[Union[str, Callable]]] = Gates.RX,
        trainable_frequencies: bool = False,
        initialization: str = "random",
        initialization_domain: List[float] = [0, 2 * jnp.pi],
        output_qubit: Union[List[int], int, None] = None,
        observables: Union[
            int, List[Union[int, List[int]]], List[op.Operation], None
        ] = None,
        shots: Optional[int] = None,
        random_seed: int = 1000,
        repeat_batch_axis: List[bool] = [True, True, True, True],
        pulse_shape: str = "gaussian",
    ) -> None:
        """
        Initialize the quantum circuit model.
        Parameters will have the shape [impl_n_layers, parameters_per_layer]
        where impl_n_layers is the number of layers provided and added by one
        depending if data_reupload is True and parameters_per_layer is given by
        the chosen ansatz.

        The model is initialized with the following parameters as defaults:
        - noise_params: None
        - execution_type: "expval"
        - shots: None

        Args:
            n_qubits (int): The number of qubits in the circuit.
            n_layers (int): The number of layers in the circuit.
            circuit_type (str, Circuit): The type of quantum circuit to use.
                If None, defaults to "no_ansatz".
            data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]],
                np.ndarray], optional):
                Whether to reupload data to the quantum device on each
                layer and qubit. Detailed re-uploading instructions can be given
                as a list/array of 0/False and 1/True with shape (n_qubits,
                n_layers) to specify where to upload the data. Defaults to True
                for applying data re-uploading to the full circuit.
            encoding (Union[str, Callable, List[str], List[Callable]], optional):
                The unitary to use for encoding the input data. Can be a string
                (e.g. "RX") or a callable (e.g. gateset.RX). Defaults to gateset.RX.
                If input is multidimensional it is assumed to be a list of
                unitaries or a list of strings.
            trainable_frequencies (bool, optional):
                Sets trainable encoding parameters for trainable frequencies.
                Defaults to False.
            initialization (str, optional): The strategy to initialize the parameters.
                Can be "random", "zeros", "zero-controlled", "pi", or "pi-controlled".
                Defaults to "random".
            output_qubit (List[int], int, optional): Deprecated alias for
                ``observables``. Forwards to ``observables`` and will be removed
                in a future release. Defaults to None.
            observables (int, List[int], List[List[int]], List[op.Operation],
                optional): Measurement specification. A qubit index, a list of
                indices, or a list of qubit groups (for $Z$-parity) selects the
                measured subsystem with the default PauliZ readout.
                Alternatively, a list of
                :class:`~jaqsi.operations.Operation` observables makes
                ``execution_type="expval"`` return one expectation value per
                observable. When None all qubits are measured. Defaults to None.
            shots (Optional[int], optional): The number of shots to use for
                the quantum device. Defaults to None.
            random_seed (int, optional): seed for the random number generator
                in initialization is "random" and for random noise parameters.
                Defaults to 1000.
            repeat_batch_axis (List[bool], optional): Each boolean in the array
                determines over which axes to parallelise computation. The axes
                correspond to [inputs, params, pulse_params, enc_pulse_params].
                Defaults to [True, True, True, True], meaning that batching is
                enabled over all axes. A 3-element list (legacy) is accepted and
                extended with a trailing True for the enc_pulse_params axis.
            pulse_shape (str, optional): Pulse envelope shape for pulse-level
                simulation. One of ``PulseEnvelope.available()``.
                Defaults to ``"gaussian"``.

        Returns:
            None
        """
        # Initialize default parameters needed for circuit evaluation
        self.n_qubits: int = n_qubits
        if output_qubit is not None:
            if observables is not None:
                raise ValueError("Pass either output_qubit or observables, not both.")
            warnings.warn(
                "output_qubit is deprecated, use observables instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            observables = output_qubit
        self.observables = observables
        self.n_layers: int = n_layers
        self.noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None
        self.shots = shots
        self.trainable_frequencies: bool = trainable_frequencies
        self.execution_type: str = "expval"
        # backward compatibility
        # TODO: consider making this more generic in future
        # (for someone wanting to control this without bothering with pulse stuff)
        if len(repeat_batch_axis) == 3:
            log.warning("Batch axis should have length 4")
            repeat_batch_axis = list(repeat_batch_axis) + [True]
        self.repeat_batch_axis: List[bool] = repeat_batch_axis

        # --- Pulse envelope ---
        pinfo.set_envelope(pulse_shape)

        # --- State Preparation ---
        try:
            self._sp = Gates.parse_gates(state_preparation, Gates)
        except ValueError as e:
            raise ValueError(f"Error parsing encodings: {e}")

        # prepare corresponding pulse parameters (always optimized pulses)
        self.sp_pulse_params = []
        for sp in self._sp:
            sp_name = sp.__name__ if hasattr(sp, "__name__") else str(sp)

            if pinfo.gate_by_name(sp_name) is not None:
                self.sp_pulse_params.append(pinfo.gate_by_name(sp_name).params)
            else:
                # gate has no pulse parametrization
                self.sp_pulse_params.append(None)

        # --- Encoding ---
        if isinstance(encoding, Encoding):
            # user wants custom strategy? do it!
            self._enc = encoding
        else:
            # use hammming encoding by default
            self._enc = Encoding("hamming", encoding)

        if self._enc.is_golomb:
            self._enc._n_qubits = n_qubits

        # Number of possible inputs
        self.n_input_feat = len(self._enc)
        log.debug(f"Number of input features: {self.n_input_feat}")

        # Trainable frequencies, default initialization as in arXiv:2309.03279v2
        self.enc_params = jnp.ones((self.n_layers, self.n_qubits, self.n_input_feat))

        # Per-feature pulse-parameter sizes/offsets used to slice
        # enc_pulse_params in _iec under "all_pulse" mode. Only encodings whose
        # gates all have a pulse parametrization are supported (golomb and
        # custom callables do not).
        # TODO: golomb should be doable but needs a closer investigation
        self._enc_pulse_sizes: List[int] = []
        self._enc_pulse_capable = not self._enc.is_golomb
        if self._enc_pulse_capable:
            for g in self._enc._gates:
                if pinfo.gate_by_name(g) is None:
                    self._enc_pulse_capable = False
                    self._enc_pulse_sizes = []
                    break
                self._enc_pulse_sizes.append(pinfo.gate_by_name(g).size)

        self._enc_pulse_offsets: List[int] = list(
            np.cumsum([0, *self._enc_pulse_sizes[:-1]])
        )
        self._enc_pulse_shape: Tuple[int, int, int] = (
            self.n_layers,
            self.n_qubits,
            sum(self._enc_pulse_sizes),
        )

        # --- Data-Reuploading ---

        # Keep as NumPy array (not JAX) so that ``if data_reupload[q, idx]``
        # in _iec remains a concrete Python bool even under jax.jit tracing.
        # note that setting this will also update self.degree and self.frequencies
        # and in consequence also self.has_dru
        self.data_reupload = data_reupload

        # check for the highest degree among all input dimensions
        if self.has_dru:
            impl_n_layers: int = n_layers + 1  # we need L+1 according to Schuld et al.
        else:
            impl_n_layers = n_layers
        log.info(f"Number of implicit layers: {impl_n_layers}.")

        # --- Ansatz ---
        # only weak check for str. We trust the user to provide sth useful
        if isinstance(circuit_type, str):
            self.pqc: Callable[[Optional[jnp.ndarray], int], int] = getattr(
                Ansaetze, circuit_type or "No_Ansatz"
            )()
        else:
            self.pqc = circuit_type()
        log.info(f"Using Ansatz {circuit_type}.")

        # calculate the shape of the parameter vector here, we will re-use this in init.
        params_per_layer = self.pqc.n_params_per_layer(self.n_qubits)
        self._params_shape: Tuple[int, int] = (impl_n_layers, params_per_layer)
        log.info(f"Parameters per layer: {params_per_layer}")

        pulse_params_per_layer = self.pqc.n_pulse_params_per_layer(self.n_qubits)
        self._pulse_params_shape: Tuple[int, int] = (
            impl_n_layers,
            pulse_params_per_layer,
        )

        # intialize to None as we can't know this yet
        self._batch_shape = None

        # this will also be re-used in the init method,
        # however, only if nothing is provided
        self._inialization_strategy = initialization
        self._initialization_domain = initialization_domain

        # ..here! where we only require a JAX random key
        self.random_key = self.initialize_params(random.key(random_seed))

        # Initializing pulse params
        self.pulse_params: jnp.ndarray = jnp.ones((1, *self._pulse_params_shape))

        log.info(f"Initialized pulse parameters with shape {self.pulse_params.shape}.")

        # Initializing encoding pulse params (element-wise scalers, ones by
        # default). Batch-first convention, mirroring pulse_params.
        self.enc_pulse_params: jnp.ndarray = jnp.ones((1, *self._enc_pulse_shape))

        log.info(
            f"Initialized encoding pulse parameters with shape "
            f"{self.enc_pulse_params.shape}."
        )

        # Initialise the jaqsi Script that wraps _variational.
        # No device selection needed - jaqsi auto-routes between statevector
        # and density-matrix simulation based on whether noise channels are
        # present on the tape.
        self.script = js.Script(f=self._variational, n_qubits=self.n_qubits)

    @property
    def noise_params(self) -> Optional[Dict[str, Union[float, Dict[str, float]]]]:
        """
        Gets the noise parameters of the model.

        Returns:
            Optional[Dict[str, float]]: A dictionary of
            noise parameters or None if not set.
        """
        return self._noise_params

    @noise_params.setter
    def noise_params(
        self, kvs: Optional[Dict[str, Union[float, Dict[str, float]]]]
    ) -> None:
        """
        Sets the noise parameters of the model.

        Typically a "noise parameter" refers to the error probability.
        ThermalRelaxation is a special case, and supports a dict as value with
        structure:
            "ThermalRelaxation":
            {
                "t1": 2000, # relative t1 time.
                "t2": 1000, # relative t2 time
                "t_factor" 1: # relative gate time factor
            },

        Args:
            kvs (Optional[Dict[str, Union[float, Dict[str, float]]]]): A
            dictionary of noise parameters. If all values are 0.0, the noise
            parameters are set to None.

        Returns:
            None
        """
        self._noise_params = self._normalize_noise_params(kvs)

    @staticmethod
    def _normalize_noise_params(
        kvs: Optional[Dict[str, Union[float, Dict[str, float]]]],
    ) -> Optional[Dict[str, Union[float, Dict[str, float]]]]:
        """
        Fill in defaults and validate a noise parameter dictionary.

        Args:
            kvs (Optional[Dict[str, Union[float, Dict[str, float]]]]): A
            dictionary of noise parameters.

        Returns:
            Optional[Dict[str, Union[float, Dict[str, float]]]]: The normalized
            dictionary, or None if all values are 0.0.
        """
        # set to None if only zero values provided
        if kvs is not None and all(v == 0.0 for v in kvs.values()):
            kvs = None

        # set default values
        if kvs is not None:
            defaults = {
                "BitFlip": 0.0,
                "PhaseFlip": 0.0,
                "Depolarizing": 0.0,
                "MultiQubitDepolarizing": 0.0,
                "AmplitudeDamping": 0.0,
                "PhaseDamping": 0.0,
                "GateError": 0.0,
                "ThermalRelaxation": None,
                "StatePreparation": 0.0,
                "Measurement": 0.0,
            }
            for key, default_val in defaults.items():
                kvs.setdefault(key, default_val)

            # check if there are any keys not supported
            for key in kvs.keys():
                if key not in defaults:
                    warnings.warn(
                        f"Noise type {key} is not supported by this package",
                        UserWarning,
                    )

            # check valid params for thermal relaxation noise channel
            tr_params = kvs["ThermalRelaxation"]
            if isinstance(tr_params, dict):
                tr_params.setdefault("t1", 0.0)
                tr_params.setdefault("t2", 0.0)
                tr_params.setdefault("t_factor", 0.0)
                valid_tr_keys = {"t1", "t2", "t_factor"}
                for k in tr_params.keys():
                    if k not in valid_tr_keys:
                        warnings.warn(
                            f"Thermal Relaxation parameter {k} is not supported "
                            f"by this package",
                            UserWarning,
                        )
                if not all(tr_params.values()) or tr_params["t2"] > 2 * tr_params["t1"]:
                    warnings.warn(
                        "Received invalid values for Thermal Relaxation noise "
                        "parameter. Thermal relaxation is not applied!",
                        UserWarning,
                    )
                    kvs["ThermalRelaxation"] = 0.0

        return kvs

    @property
    def output_qubit(self) -> List[int]:
        """Deprecated alias for :attr:`observables`; returns the measured wires."""
        warnings.warn(
            "output_qubit is deprecated, use observables instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self._measured_wires

    @output_qubit.setter
    def output_qubit(self, value: Union[int, List[int]]) -> None:
        warnings.warn(
            "output_qubit is deprecated, use observables instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        self.observables = value

    @property
    def observables(self) -> List:
        """The custom :class:`~jaqsi.operations.Operation` observables,
        or the list of measured wires when using the default PauliZ readout.

        With a list of observables, ``__call__`` and ``execution_type="expval"``
        returns one expectation value per observable instead of one ``PauliZ``
        per measured qubit.
        """
        return (
            self._observables if self._observables is not None else self._measured_wires
        )

    @observables.setter
    def observables(self, value: Union[int, List, None]) -> None:
        if value is None:
            self._observables = None
            self._measured_wires = list(range(self.n_qubits))
        elif (
            isinstance(value, list)
            and value
            and all(isinstance(o, op.Operation) for o in value)
        ):
            self._observables = list(value)
            self._measured_wires = list(range(self.n_qubits))
        elif isinstance(value, list) and any(
            isinstance(o, op.Operation) for o in value
        ):
            raise ValueError(
                "observables list must contain either qubit indices or "
                "Operation objects, not a mix."
            )
        else:
            # qubit specification: normalize into the measured wire list
            self._observables = None
            if isinstance(value, list):
                assert len(value) <= self.n_qubits, (
                    f"Size of observables {len(value)} cannot be larger than "
                    f"number of qubits {self.n_qubits}."
                )
                self._measured_wires = value
            elif isinstance(value, int):
                if value == -1:
                    self._measured_wires = list(range(self.n_qubits))
                else:
                    assert value < self.n_qubits, (
                        f"Output qubit {value} cannot be larger than {self.n_qubits}."
                    )
                    self._measured_wires = [value]
            else:
                self._measured_wires = value

        # recompute the result shape for the (possibly new) observable count
        if hasattr(self, "_execution_type"):
            self.execution_type = self.execution_type

    @property
    def execution_type(self) -> str:
        """
        Gets the execution type of the model.

        Returns:
            str: The execution type, one of 'density', 'expval', or 'probs'.
        """
        return self._execution_type

    @execution_type.setter
    def execution_type(self, value: str) -> None:
        self._result_shape = self._compute_result_shape(value)

        if value == "state" and not self.all_qubit_measurement:
            warnings.warn(
                f"{value} measurement ignores the measured subsystem, which is "
                f"{self._measured_wires}.",
                UserWarning,
            )

        if value != "expval" and getattr(self, "_observables", None) is not None:
            warnings.warn(
                f"Custom observables are ignored for execution_type={value!r}.",
                UserWarning,
            )

        if value == "probs" and self.shots is None:
            warnings.warn(
                "Setting execution_type to probs without specifying shots.",
                UserWarning,
            )

        if value == "density" and self.shots is not None:
            raise ValueError("Setting execution_type to density with shots not None.")

        self._execution_type = value

    def _compute_result_shape(self, execution_type: str) -> Tuple[int, ...]:
        """
        Derive the per-sample output shape for an execution type.

        Args:
            execution_type (str): One of "density", "expval", "probs", "state".

        Returns:
            Tuple[int, ...]: The output shape of a single sample.

        Raises:
            ValueError: If execution_type is not supported.
        """
        if execution_type == "density":
            return (
                2 ** len(self._measured_wires),
                2 ** len(self._measured_wires),
            )
        elif execution_type == "expval":
            # custom observables (if provided) fix the number of expectation
            # values; otherwise one PauliZ (or Z-parity) per measured qubit.
            if getattr(self, "_observables", None) is not None:
                return (len(self._observables),)
            else:
                return (len(self._measured_wires),)
        elif execution_type == "probs":
            # in case this is a list of parities,
            # each pair has 2^len(qubits) probabilities
            return (
                (2,) * len(self._measured_wires)
                if isinstance(self._measured_wires, (Tuple, List))
                else (2,)
            )
        elif execution_type == "state":
            return (2 ** len(self._measured_wires),)
        else:
            raise ValueError(f"Invalid execution type: {execution_type}.")

    @property
    def shots(self) -> Optional[int]:
        """
        Gets the number of shots to use for the quantum device.

        Returns:
            Optional[int]: The number of shots.
        """
        return self._shots

    @shots.setter
    def shots(self, value: Optional[int]) -> None:
        """
        Sets the number of shots to use for the quantum device.

        Args:
            value (Optional[int]): The number of shots.
            If an integer less than or equal to 0 is provided, it is set to None.

        Returns:
            None
        """
        if type(value) is int and value <= 0:
            value = None
        self._shots = value

    @property
    def params(self) -> jnp.ndarray:
        """Get the variational parameters of the model."""
        return self._params

    @params.setter
    def params(self, value: jnp.ndarray) -> None:
        """Set the variational parameters, ensuring batch dimension exists."""
        if len(value.shape) == 2:
            value = value.reshape(1, *value.shape)

        self._params = value

    @property
    def enc_params(self) -> jnp.ndarray:
        """Get the encoding parameters used for input transformation."""
        return self._enc_params

    @enc_params.setter
    def enc_params(self, value: jnp.ndarray) -> None:
        """Set the encoding parameters."""
        self._enc_params = value

    @property
    def pulse_params(self) -> jnp.ndarray:
        """Get the pulse parameters for pulse-mode gate execution."""
        return self._pulse_params

    @pulse_params.setter
    def pulse_params(self, value: jnp.ndarray) -> None:
        """Set the pulse parameters."""
        self._pulse_params = value

    @property
    def enc_pulse_params(self) -> jnp.ndarray:
        """Get the encoding pulse parameters for all_pulse-mode execution."""
        return self._enc_pulse_params

    @enc_pulse_params.setter
    def enc_pulse_params(self, value: jnp.ndarray) -> None:
        """Set the encoding pulse parameters."""
        self._enc_pulse_params = value

    @property
    def data_reupload(self) -> np.ndarray:
        """Get the data reupload mask."""
        return self._data_reupload

    @data_reupload.setter
    def data_reupload(
        self,
        value: Union[bool, List[List[bool]], List[List[List[bool]]], np.ndarray],
    ) -> None:
        """Set the data reupload mask.

        Always converts to a concrete NumPy boolean array so that
        ``if data_reupload[q, idx]`` in :meth:`_iec` remains a plain
        Python ``bool`` even inside JAX-traced functions (jit / grad / vmap).
        """
        # Process data reuploading strategy and set degree
        if not isinstance(value, bool):
            if not isinstance(value, np.ndarray):
                value = np.array(value)

            if len(value.shape) == 2:
                assert value.shape == (
                    self.n_layers,
                    self.n_qubits,
                ), (
                    f"Data reuploading array has wrong shape. \
                    Expected {(self.n_layers, self.n_qubits)} or\
                    {(self.n_layers, self.n_qubits, self.n_input_feat)},\
                    got {value.shape}."
                )
                value = value.reshape(*value.shape, 1)
                value = np.repeat(value, self.n_input_feat, axis=2)

            assert value.shape == (
                self.n_layers,
                self.n_qubits,
                self.n_input_feat,
            ), (
                f"Data reuploading array has wrong shape. \
                Expected {(self.n_layers, self.n_qubits, self.n_input_feat)},\
                got {value.shape}."
            )

            log.debug(f"Data reuploading array:\n{value}")
        else:
            if value:
                value = np.ones((self.n_layers, self.n_qubits, self.n_input_feat))
                log.debug("Full data reuploading.")
            else:
                value = np.zeros((self.n_layers, self.n_qubits, self.n_input_feat))
                value[0][0] = 1
                log.debug("No data reuploading.")

        # convert to boolean values
        self._data_reupload = np.asarray(value).astype(bool)

        self.degree: Tuple = tuple(
            self._enc.get_n_freqs(self.data_reupload[..., i])
            for i in range(self.n_input_feat)
        )

        self.frequencies: Tuple = tuple(
            self._enc.get_spectrum(self.data_reupload[..., i])
            for i in range(self.n_input_feat)
        )

        # Cache has_dru as a plain Python bool so that it can be used in
        # Python ``if`` statements even inside JAX-traced functions.
        self._has_dru: bool = bool(max(int(np.max(f)) for f in self._frequencies) > 1)

    @property
    def degree(self) -> Tuple:
        """Get the degree of the model."""
        return self._degree

    @degree.setter
    def degree(self, value: Tuple):
        self._degree = value

    @property
    def frequencies(self) -> Tuple:
        """Get the frequencies of the model."""
        return self._frequencies

    @frequencies.setter
    def frequencies(self, value: Tuple):
        self._frequencies = value

    def exact_spectrum(self, method: str = "tree") -> Tuple[np.ndarray, ...]:
        """Compute the exact per-feature Fourier spectrum via the FourierTree.

        Unlike :attr:`frequencies` -- a naive per-feature estimate derived purely
        from the encoding, which can *overestimate* the spectrum (some
        coefficients are constrained to zero for all parameters) -- this builds
        the analytical Fourier tree (Nemkov et al.) and returns, for each input
        feature, the integer frequencies whose Fourier coefficient is not
        identically zero.  The result is always a subset of :attr:`frequencies`.

        The support is derived purely symbolically (no parameter sampling): see
        :meth:`~qml_essentials.coefficients.FourierTree.get_exact_support`.
        With ``method="tree"`` (default), frequencies whose contributions cancel
        identically across tree paths (e.g. two consecutive encodings combining
        into a single rotation) are excluded exactly; this enumerates the
        explicit tree, which can be infeasible for deep entangling circuits.
        With ``method="dp"``, a merged-state dynamic program derives the support
        without enumerating paths, which scales to deep circuits at the cost of
        not detecting identical cross-path cancellations.

        Requires a Clifford + Pauli-rotation ansatz (see
        :class:`~qml_essentials.pauli.PauliCircuit`); other gate sets raise
        ``NotImplementedError`` during tree construction.

        Args:
            method (str): ``"tree"`` (fully exact) or ``"dp"`` (scalable).

        Returns:
            Tuple[np.ndarray, ...]: One sorted integer frequency array per input
            feature (same layout as :attr:`frequencies`).
        """
        from qml_essentials.coefficients import FourierTree  # avoid circular imp.

        tree = FourierTree(self)

        # Position of each model feature within the tree's frequency vectors.
        feature_pos = {feat: i for i, feat in enumerate(tree.features)}

        # Union of the symbolic supports over all observables (roots).
        support = set()
        for freqs in tree.get_exact_support(method=method):
            farr = np.asarray(freqs)
            for k in range(farr.shape[0]):
                key = (
                    (int(farr[k]),)
                    if farr.ndim == 1
                    else tuple(int(v) for v in farr[k])
                )
                support.add(key)

        spectrum = []
        for feat in range(self.n_input_feat):
            if support and feat in feature_pos:
                pos = feature_pos[feat]
                vals = sorted({k[pos] for k in support})
            else:
                vals = [0]
            spectrum.append(np.array(vals, dtype=int))
        return tuple(spectrum)

    @property
    def has_dru(self) -> bool:
        """Check if the model has data reupload."""
        return self._has_dru

    @property
    def all_qubit_measurement(self) -> bool:
        """Check if measurement is performed on all qubits."""
        return self._measured_wires == list(range(self.n_qubits))

    @property
    def batch_shape(self) -> Tuple[int, ...]:
        """
        Get the batch shape (B_I, B_P, B_R, B_E).
        If the model was not called before,
        it returns (1, 1, 1, 1).

        Returns:
            Tuple[int, ...]: Tuple of (input_batch, param_batch, pulse_batch,
                enc_pulse_batch). Returns (1, 1, 1, 1) if model has not been
                called yet.
        """
        if self._batch_shape is None:
            log.debug("Model was not called yet. Returning (1,1,1,1) as batch shape.")
            return (1, 1, 1, 1)
        return self._batch_shape

    @property
    def eff_batch_shape(self) -> Tuple[int, ...]:
        """
        Get the effective batch shape after applying repeat_batch_axis mask.

        Returns:
            Tuple[int, ...]: Effective batch dimensions, excluding zeros.
        """
        return self._eff_batch_shape_of(self.batch_shape)

    def _eff_batch_shape_of(self, batch_shape: Tuple[int, ...]) -> Tuple[int, ...]:
        """
        Apply the repeat_batch_axis mask to a given batch shape.

        Args:
            batch_shape (Tuple[int, ...]): Batch shape (B_I, B_P, B_R, B_E).

        Returns:
            Tuple[int, ...]: Effective batch dimensions, excluding zeros.
        """
        batch_shape = np.array(batch_shape) * self.repeat_batch_axis
        return batch_shape[batch_shape != 0]

    def initialize_params(
        self,
        random_key: Optional[random.PRNGKey] = None,
        repeat: int = 1,
        initialization: Optional[str] = None,
        initialization_domain: Optional[List[float]] = None,
    ) -> random.PRNGKey:
        """
        Initialize the variational parameters of the model.

        Args:
            random_key (Optional[random.PRNGKey]): JAX random key for initialization.
                If None, uses the model's internal random key.
            repeat (int): Number of parameter sets to create (batch dimension).
                Defaults to 1.
            initialization (Optional[str]): Strategy for parameter initialization.
                Options: "random", "zeros", "pi", "zero-controlled", "pi-controlled".
                If None, uses the strategy specified in the constructor.
            initialization_domain (Optional[List[float]]): Domain [min, max] for
                random initialization. If None, uses the domain from constructor.

        Returns:
            random.PRNGKey: Updated random key after initialization.

        Raises:
            Exception: If an invalid initialization method is specified.
        """
        # Initializing params
        params_shape = (repeat, *self._params_shape)

        # use existing strategy if not specified
        initialization = initialization or self._inialization_strategy
        initialization_domain = initialization_domain or self._initialization_domain

        random_key, sub_key = safe_random_split(
            random_key if random_key is not None else self.random_key
        )

        def set_control_params(params: jnp.ndarray, value: float) -> jnp.ndarray:
            indices = self.pqc.get_control_indices(self.n_qubits)
            if indices is None:
                warnings.warn(
                    f"Specified {initialization} but circuit\
                    does not contain controlled rotation gates.\
                    Parameters are intialized randomly.",
                    UserWarning,
                )
            else:
                np_params = np.array(params)
                np_params[:, :, indices[0] : indices[1] : indices[2]] = (
                    np.ones_like(params[:, :, indices[0] : indices[1] : indices[2]])
                    * value
                )
                params = jnp.array(np_params)
            return params

        if initialization == "random":
            self.params: jnp.ndarray = random.uniform(
                sub_key,
                params_shape,
                minval=initialization_domain[0],
                maxval=initialization_domain[1],
            )
        elif initialization == "zeros":
            self.params: jnp.ndarray = jnp.zeros(params_shape)
        elif initialization == "pi":
            self.params: jnp.ndarray = jnp.ones(params_shape) * jnp.pi
        elif initialization == "zero-controlled":
            self.params: jnp.ndarray = random.uniform(
                sub_key,
                params_shape,
                minval=initialization_domain[0],
                maxval=initialization_domain[1],
            )
            self.params = set_control_params(self.params, 0)
        elif initialization == "pi-controlled":
            self.params: jnp.ndarray = random.uniform(
                sub_key,
                params_shape,
                minval=initialization_domain[0],
                maxval=initialization_domain[1],
            )
            self.params = set_control_params(self.params, jnp.pi)
        else:
            raise Exception("Invalid initialization method")

        log.info(
            f"Initialized parameters with shape {self.params.shape}\
            using strategy {initialization}."
        )

        return random_key

    def next_key(self) -> random.PRNGKey:
        """
        Advance the internal random key and return a fresh sub key.

        Intended for stochastic execution inside a JAX transform: a jitted
        call is traced once and replays the key that was current at trace
        time, so fresh randomness has to enter as an argument. Call this
        outside the transform and pass the result as ``random_key``. Since the
        key is an argument rather than a constant, this does not trigger
        recompilation.

        Returns:
            random.PRNGKey: Fresh sub key, split off the internal key.
        """
        self.random_key, sub_key = safe_random_split(self.random_key)
        return sub_key

    def transform_input(
        self, inputs: jnp.ndarray, enc_params: jnp.ndarray
    ) -> jnp.ndarray:
        """
        Transform input data by scaling with encoding parameters.

        Implements the input transformation as described in arXiv:2309.03279v2,
        where inputs are linearly scaled by encoding parameters before being
        used in the quantum circuit.

        Args:
            inputs (jnp.ndarray): Input data point of shape (n_input_feat,) or
                (batch_size, n_input_feat).
            enc_params (jnp.ndarray): Encoding weight scalar or vector used to
                scale the input.

        Returns:
            jnp.ndarray: Transformed input, element-wise product of inputs
                and enc_params.
        """
        return inputs * enc_params

    def _iec(
        self,
        inputs: jnp.ndarray,
        data_reupload: np.ndarray,
        enc: Encoding,
        enc_params: jnp.ndarray,
        noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
        random_key: Optional[random.PRNGKey] = None,
        enc_pulse_params: Optional[jnp.ndarray] = None,
        pulse: bool = False,
    ) -> None:
        """
        Apply Input Encoding Circuit (IEC) with angle encoding.

        Encodes classical input data into the quantum circuit using rotation
        gates (e.g., RX, RY, RZ). Supports data re-uploading at specified
        positions in the circuit.

        For Golomb encoding, a single multi-qubit diagonal unitary is applied
        to all qubits simultaneously instead of per-qubit rotation gates.

        Args:
            inputs (jnp.ndarray): Input data of shape (n_input_feat,) or
                (batch_size, n_input_feat).
            data_reupload (np.ndarray): Boolean array of shape (n_qubits, n_input_feat)
                indicating where to apply encoding gates.
            enc (Encoding): Encoding strategy containing the encoding gate functions.
            enc_params (jnp.ndarray): Encoding parameters of shape
                (n_qubits, n_input_feat) used to scale inputs.
            noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
                Noise parameters for gate-level noise simulation. Defaults to None.
            random_key (Optional[random.PRNGKey]): JAX random key for stochastic
                noise. Defaults to None.
            enc_pulse_params (Optional[jnp.ndarray]): Encoding pulse-parameter
                scalers of shape (n_qubits, n_enc_pulse_per_qubit) for the
                current layer. Used when the encoding gates run at pulse level,
                i.e. the model-level mode is "enc_pulse" or "all_pulse".
                Defaults to None.
            pulse (bool): Whether the encoding gates run at pulse level.
                This is the backend selected for the encoding group, distinct
                from the model-level modes (unitary, ansatz_pulse, enc_pulse,
                all_pulse). Defaults to False.

        Returns:
            None: Gates are applied in-place to the quantum circuit.
        """
        # --- Golomb encoding: single multi-qubit gate on all qubits --------
        if enc.is_golomb:
            idx = 0  # Golomb encoding supports a single input feature
            # Check if any qubit has re-uploading enabled for this layer
            if data_reupload[:, idx].any():
                random_key, sub_key = safe_random_split(random_key)
                # Use the mean of enc_params across qubits as scalar scaling
                # (Golomb acts on all qubits jointly)
                mean_enc_param = jnp.mean(enc_params[:, idx])
                all_wires = list(range(self.n_qubits))
                enc[idx](
                    self.transform_input(inputs[..., idx], mean_enc_param),
                    wires=all_wires,
                    noise_params=noise_params,
                    random_key=sub_key,
                )
            return

        # --- Standard per-qubit encoding -----------------------------------
        for q in range(self.n_qubits):
            # use the last dimension of the inputs (feature dimension)
            for idx in range(inputs.shape[-1]):
                if data_reupload[q, idx]:
                    random_key, sub_key = safe_random_split(random_key)
                    # TODO: consider merging this with the pulses.py manager
                    pulse_kwargs = {}
                    if pulse:
                        # scale the calibrated pulse params by this gate's
                        # scalers, as the pulse manager does for the ansatz
                        off = self._enc_pulse_offsets[idx]
                        size = self._enc_pulse_sizes[idx]
                        base = pinfo.gate_by_name(enc._gates[idx]).params
                        pulse_kwargs = dict(
                            pulse_params=base * enc_pulse_params[q, off : off + size],
                            pulse=True,
                        )

                    # use elipsis to index only the last dimension
                    # as inputs are generally *not* qubit dependent
                    enc[idx](
                        self.transform_input(inputs[..., idx], enc_params[q, idx]),
                        wires=q,
                        noise_params=noise_params,
                        random_key=sub_key,
                        **pulse_kwargs,
                    )

    @staticmethod
    def _debatch(value: jnp.ndarray, ndim: int) -> jnp.ndarray:
        """
        Drop a leading singleton batch axis (batch-first convention).

        Args:
            value (jnp.ndarray): Array to de-batch.
            ndim (int): Rank of a single (un-batched) element.

        Returns:
            jnp.ndarray: The array without its leading axis if that axis is a
                singleton batch dimension, otherwise the array unchanged.
        """
        if len(value.shape) > ndim and value.shape[0] == 1:
            return value[0]
        return value

    def _self_fallback(self, value: Any, name: str, warn: bool) -> Any:
        """
        Fall back to the model's own attribute when a parameter is not given.

        Args:
            value (Any): The provided value, or None.
            name (str): Name of the attribute to fall back to.
            warn (bool): Whether to warn when the fallback is used.

        Returns:
            Any: The provided value, or ``self.<name>`` if value is None.
        """
        if value is not None:
            return value
        if warn:
            warnings.warn(
                "Explicit call to `_circuit` or `_variational` detected: "
                f"`{name}` is None, using `self.{name}` instead.",
                RuntimeWarning,
            )
        return getattr(self, name)

    def _variational(
        self,
        params: jnp.ndarray,
        inputs: jnp.ndarray,
        pulse_params: Optional[jnp.ndarray] = None,
        random_key: Optional[random.PRNGKey] = None,
        enc_params: Optional[jnp.ndarray] = None,
        enc_pulse_params: Optional[jnp.ndarray] = None,
        gate_mode: str = "unitary",
        noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
    ) -> None:
        """
        Build the variational quantum circuit structure.

        Constructs the circuit by applying state preparation, alternating
        variational ansatz layers with input encoding layers, and optional
        noise channels.

        The first six parameters (after ``self``) - ``params``, ``inputs``,
        ``pulse_params``, ``random_key``, ``enc_params``, ``enc_pulse_params`` -
        are the batchable positional arguments.
        The remaining keyword arguments are broadcast across the batch.

        Args:
            params (jnp.ndarray): Variational parameters of shape
                (n_layers, n_params_per_layer).
            inputs (jnp.ndarray): Input data of shape (n_input_feat,).
            pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers of shape
                (n_layers, n_pulse_params_per_layer) for pulse-mode execution.
                Defaults to None (uses model's pulse_params).
            random_key (Optional[random.PRNGKey]): JAX random key for stochastic
                operations. Defaults to None.
            enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
                (n_qubits, n_input_feat). Defaults to None (uses model's enc_params).
            enc_pulse_params (Optional[jnp.ndarray]): Encoding pulse-parameter
                scalers of shape (n_layers, n_qubits, n_enc_pulse_per_qubit) for
                "all_pulse" execution. Defaults to None (uses model's
                enc_pulse_params).
            gate_mode (str): Gate execution mode, one of "unitary",
                "ansatz_pulse", "enc_pulse" or "all_pulse". "ansatz_pulse" runs
                the ansatz and state preparation as pulses (encoding stays
                unitary); "enc_pulse" runs only the encoding gates as pulses;
                "all_pulse" runs both as pulses. Defaults to "unitary".
            noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
                Noise parameters for simulation. Defaults to None.

        Returns:
            None: Gates are applied in-place to the quantum circuit.

        Note:
            Issues RuntimeWarning if called directly without providing parameters
            that would normally be passed through the forward method.
        """
        # which backend the ansatz / state-prep gates and the encoding gates use
        use_pulse, enc_use_pulse = GATE_MODES[gate_mode]

        # TODO: rework and double check params shape
        params = self._debatch(params, 2)
        inputs = self._debatch(inputs, 1)

        # TODO: Raise warning if trainable frequencies is True, or similar. I.e., no
        #   warning if user does not care for frequencies or enc_params
        enc_params = self._self_fallback(
            enc_params, "enc_params", self.trainable_frequencies
        )

        pulse_params = self._self_fallback(pulse_params, "pulse_params", use_pulse)
        pulse_params = self._debatch(pulse_params, 2)

        enc_pulse_params = self._self_fallback(
            enc_pulse_params, "enc_pulse_params", enc_use_pulse
        )
        enc_pulse_params = self._debatch(enc_pulse_params, 3)

        noise_params = self._self_fallback(
            noise_params, "noise_params", self.noise_params is not None
        )

        if noise_params is not None:
            random_key = self._self_fallback(random_key, "random_key", True)
            self._apply_state_prep_noise(noise_params=noise_params)

        # state preparation
        for q in range(self.n_qubits):
            for _sp, sp_pulse_params in zip(self._sp, self.sp_pulse_params):
                random_key, sub_key = safe_random_split(random_key)
                _sp(
                    wires=q,
                    pulse_params=sp_pulse_params,
                    noise_params=noise_params,
                    random_key=sub_key,
                    pulse=use_pulse,
                )

        # circuit building
        for layer in range(0, self.n_layers):
            random_key, sub_key = safe_random_split(random_key)
            # ansatz layers
            self.pqc(
                params[layer],
                self.n_qubits,
                pulse_params=pulse_params[layer],
                noise_params=noise_params,
                random_key=sub_key,
                pulse=use_pulse,
            )

            random_key, sub_key = safe_random_split(random_key)
            # encoding layers
            self._iec(
                inputs,
                data_reupload=self.data_reupload[layer],
                enc=self._enc,
                enc_params=enc_params[layer],
                noise_params=noise_params,
                random_key=sub_key,
                enc_pulse_params=enc_pulse_params[layer],
                pulse=enc_use_pulse,
            )

        # final ansatz layer
        if self.has_dru:  # same check as in init
            random_key, sub_key = safe_random_split(random_key)
            self.pqc(
                params[self.n_layers],
                self.n_qubits,
                pulse_params=pulse_params[-1],
                noise_params=noise_params,
                random_key=sub_key,
                pulse=use_pulse,
            )

        # channel noise
        if noise_params is not None:
            self._apply_general_noise(noise_params=noise_params)

    def _build_obs(
        self, execution_type: Optional[str] = None
    ) -> Tuple[str, List[op.Operation]]:
        """Build the jaqsi measurement type and observable list.

        Translates the model's ``execution_type`` and ``observables``
        settings into parameters suitable for
        :meth:`~jaqsi.Script.execute`.

        Args:
            execution_type: Measurement type to build for.  If ``None``, the
                model's current ``execution_type`` is used.

        Returns:
            Tuple ``(meas_type, obs)`` where *meas_type* is one of
            ``"expval"``, ``"probs"``, ``"density"``, ``"state"`` and *obs*
            is a (possibly empty) list of :class:`Operation` observables.
        """
        if execution_type is None:
            execution_type = self.execution_type

        if execution_type == "density":
            return "density", []

        if execution_type == "state":
            return "state", []

        if execution_type == "expval":
            if self._observables is not None:
                return "expval", list(self._observables)
            obs: List[op.Operation] = []
            for qubit_spec in self._measured_wires:
                if isinstance(qubit_spec, int):
                    obs.append(gateset.PauliZ(wires=qubit_spec))
                else:
                    # parity: Z \\otimes Z \\otimes …
                    obs.append(js.build_parity_observable(list(qubit_spec)))
            return "expval", obs

        if execution_type == "probs":
            # probs are computed on the full system; subsystem
            # marginalisation is handled in _postprocess_res
            return "probs", []

        raise ValueError(f"Invalid execution_type: {execution_type}.")

    def _apply_state_prep_noise(
        self, noise_params: Dict[str, Union[float, Dict[str, float]]]
    ) -> None:
        """
        Apply state preparation noise to all qubits.

        Simulates imperfect state preparation by applying BitFlip errors
        to each qubit with the specified probability.

        Args:
            noise_params (Dict[str, Union[float, Dict[str, float]]]): Dictionary
                containing noise parameters. Uses the "StatePreparation" key
                for the BitFlip probability.

        Returns:
            None: Noise channels are applied in-place to the circuit.
        """
        p = noise_params.get("StatePreparation", 0.0)
        if p > 0:
            for q in range(self.n_qubits):
                noise.BitFlip(p, wires=q)

    def _apply_general_noise(
        self, noise_params: Dict[str, Union[float, Dict[str, float]]]
    ) -> None:
        """
        Apply general noise channels to all qubits.

        Applies various decoherence and error channels after the circuit
        execution, simulating environmental noise effects.

        Args:
            noise_params (Dict[str, Union[float, Dict[str, float]]]): Dictionary
                containing noise parameters with the following supported keys:
                - "AmplitudeDamping" (float): Probability for amplitude damping.
                - "PhaseDamping" (float): Probability for phase damping.
                - "Measurement" (float): Probability for measurement error (BitFlip).
                - "ThermalRelaxation" (Dict): Dictionary with keys "t1", "t2",
                  "t_factor" for thermal relaxation simulation.

        Returns:
            None: Noise channels are applied in-place to the circuit.

        Note:
            Gate-level noise (e.g., GateError) is handled separately in the
            Gates.Noise module and applied at the individual gate level.
        """
        amp_damp = noise_params.get("AmplitudeDamping", 0.0)
        phase_damp = noise_params.get("PhaseDamping", 0.0)
        thermal_relax = noise_params.get("ThermalRelaxation", 0.0)
        meas = noise_params.get("Measurement", 0.0)
        for q in range(self.n_qubits):
            if amp_damp > 0:
                noise.AmplitudeDamping(amp_damp, wires=q)
            if phase_damp > 0:
                noise.PhaseDamping(phase_damp, wires=q)
            if meas > 0:
                noise.BitFlip(meas, wires=q)
            if isinstance(thermal_relax, dict):
                t1 = thermal_relax["t1"]
                t2 = thermal_relax["t2"]
                t_factor = thermal_relax["t_factor"]
                circuit_depth = self._get_circuit_depth()
                tg = circuit_depth * t_factor
                noise.ThermalRelaxationError(1.0, t1, t2, tg, q)

    def _get_circuit_depth(self, inputs: Optional[jnp.ndarray] = None) -> int:
        """
        Calculate the depth of the quantum circuit.

        Records the circuit onto a tape (without noise) and computes the
        depth as the length of the critical path: each gate is scheduled
        at the earliest time step after all of its qubits are free.

        Args:
            inputs (Optional[jnp.ndarray]): Input data for circuit evaluation.
                If None, default zero inputs are used.

        Returns:
            int: The circuit depth (longest path of gates in the circuit).
        """
        # Return cached value if available
        if hasattr(self, "_cached_circuit_depth"):
            return self._cached_circuit_depth

        inputs = self._inputs_validation(inputs)

        # Temporarily clear noise_params to prevent _variational from
        # picking them up (which would call _apply_general_noise ->
        # _get_circuit_depth again, causing infinite recursion).
        saved_noise = self._noise_params
        self._noise_params = None

        with recording() as tape:
            self._variational(
                self.params[0] if self.params.ndim == 3 else self.params,
                inputs[0] if inputs.ndim == 2 else inputs,
                noise_params=None,
            )

        self._noise_params = saved_noise

        # Filter out noise channels - only count unitary gates
        ops = [o for o in tape if not isinstance(o, KrausChannel)]

        if not ops:
            self._cached_circuit_depth = 0
            return 0

        # Schedule each gate at the earliest time step where all its wires
        # are free.  ``wire_busy[q]`` tracks the next free time step for
        # qubit ``q``.
        wire_busy: Dict[int, int] = {}
        depth = 0
        for gate in ops:
            start = max((wire_busy.get(w, 0) for w in gate.wires), default=0)
            end = start + 1
            for w in gate.wires:
                wire_busy[w] = end
            depth = max(depth, end)

        self._cached_circuit_depth = depth
        return depth

    def draw(
        self,
        inputs: Optional[jnp.ndarray] = None,
        figure: str = "text",
        **kwargs: Any,
    ) -> Union[str, Any]:
        """Visualize the quantum circuit.

        Records the circuit tape (without noise) and renders the gate
        sequence using the requested backend.

        Args:
            inputs (Optional[jnp.ndarray]): Input data for the circuit.
                If ``None``, default zero inputs are used.
            figure (str): Rendering backend.  One of:

                * ``"text"``  - ASCII art (returned as a ``str``).
                * ``"mpl"``   - Matplotlib figure (returns ``(fig, ax)``).
                * ``"tikz"``  - LaTeX/TikZ ``quantikz`` code (returns a
                  :class:`TikzFigure`).
                * ``"pulse"`` - Pulse schedule (returns ``(fig, axes)``).
                  Only meaningful for pulse-mode models.

            **kwargs: Extra options forwarded to the drawing backend
                (e.g. ``gate_values=True``).

        Returns:
            Depends on figure:

            * ``"text"``  -> ``str``
            * ``"mpl"``   -> ``(matplotlib.figure.Figure, matplotlib.axes.Axes)``
            * ``"tikz"``  -> :class:`TikzFigure`

        Raises:
            ValueError: If figure is not one of the supported modes.
        """
        inputs = self._inputs_validation(inputs)
        params = self.params[0] if self.params.ndim == 3 else self.params
        inp = inputs[0] if inputs.ndim == 2 else inputs

        if figure == "pulse":
            return self.draw_pulse(inputs=inputs, **kwargs)

        # Record without noise to get a clean circuit
        saved_noise = self._noise_params
        self._noise_params = None

        draw_script = js.Script(f=self._variational, n_qubits=self.n_qubits)
        result = draw_script.draw(
            figure=figure,
            args=(params, inp),
            kwargs={"noise_params": None},
            **kwargs,
        )

        self._noise_params = saved_noise
        return result

    def draw_pulse(
        self,
        inputs: Optional[jnp.ndarray] = None,
        **kwargs: Any,
    ) -> Any:
        """Visualize the pulse schedule for the circuit.

        Records the circuit in pulse mode and collects PulseEvents
        automatically via the pulse-event tape, then renders them.

        State preparation, ansatz and encoding gates are all rendered as
        pulses. Encodings without a pulse parametrization (golomb and custom
        callables) are omitted from the schedule.

        Args:
            inputs: Input data.  If ``None``, default zero inputs are used.
            **kwargs: Forwarded to
                :func:`~jaqsi.drawing.draw_pulse_schedule`
                (e.g. ``show_carrier=True``, ``n_samples=300``).

        Returns:
            ``(fig, axes)`` — Matplotlib Figure and array of Axes.
        """
        if "gate_mode" in kwargs:
            warnings.warn(
                "draw_pulse no longer takes gate_mode, every gate group with a "
                "pulse representation is drawn.",
                DeprecationWarning,
                stacklevel=2,
            )
            kwargs.pop("gate_mode")

        inputs = self._inputs_validation(inputs)
        params = self.params[0] if self.params.ndim == 3 else self.params
        inp = inputs[0] if inputs.ndim == 2 else inputs

        # pass the model's own pulse parameters, so that _variational does not
        # fall back to them with a warning. Both are batch-first, so drawing
        # picks the first set, same as params above
        record_kwargs: Dict[str, Any] = {
            "gate_mode": "all_pulse" if self._enc_pulse_capable else "ansatz_pulse",
            "noise_params": None,
            "pulse_params": self.pulse_params[0],
        }
        if self._enc_pulse_capable:
            record_kwargs["enc_pulse_params"] = self.enc_pulse_params[0]

        draw_script = js.Script(f=self._variational, n_qubits=self.n_qubits)
        return draw_script.draw(
            figure="pulse",
            args=(params, inp),
            kwargs=record_kwargs,
            **kwargs,
        )

    def __repr__(self) -> str:
        """Return text representation of the quantum circuit model."""
        return self.draw(figure="text")

    def __str__(self) -> str:
        """Return string representation of the quantum circuit model."""
        return self.draw(figure="text")

    def _params_validation(
        self, params: Optional[jnp.ndarray], stash: bool = True
    ) -> jnp.ndarray:
        """
        Validate and normalize variational parameters.

        Ensures parameters have the correct shape with a batch dimension,
        and updates the model's internal parameters if new ones are provided.

        Args:
            params (Optional[jnp.ndarray]): Variational parameters to validate.
                If None, returns the model's current parameters.
            stash (bool): Whether to store the validated parameters on the
                model. Set to False on the pure :meth:`apply` path, where
                stashing a JAX tracer would leak it across calls. Defaults
                to True.

        Returns:
            jnp.ndarray: Validated parameters with shape
                (batch_size, n_layers, n_params_per_layer).
        """
        # append batch axis if not provided
        if params is not None:
            if len(params.shape) == 2:
                # jnp (not np) so params stays a JAX array under autodiff /
                # jit; mirrors the pulse_params handling below.
                params = jnp.expand_dims(params, axis=0)

            # Avoid stashing JAX tracers on ``self``: under an outer
            # transform (e.g. ``jit``/``jacrev``) the tracer becomes invalid
            # once the transform returns, and a subsequent read of
            # ``self.params`` would feed a leaked tracer into the next
            # call (raising ``UnexpectedTracerError``).
            if stash and not isinstance(params, jax.core.Tracer):
                self.params = params
            elif stash:
                log.debug(
                    "`params` is a JAX tracer; `self.params` is left at its "
                    "previous value. Anything reading model state afterwards "
                    "(draw, Entanglement, Expressibility, or a call that omits "
                    "`params`) will see the stale parameters - assign "
                    "`model.params` explicitly if you need the state to follow."
                )
        else:
            params = self.params

        return params

    def _pulse_params_validation(
        self, pulse_params: Optional[jnp.ndarray], stash: bool = True
    ) -> jnp.ndarray:
        """
        Validate and normalize pulse parameters.

        Ensures pulse parameters are set, using model defaults if not provided.

        Args:
            pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers.
                If None, returns the model's current pulse parameters.
            stash (bool): Whether to store the validated parameters on the
                model. See :meth:`_params_validation`. Defaults to True.

        Returns:
            jnp.ndarray: Validated pulse parameters with shape
                (batch_size, n_layers, n_pulse_params_per_layer).
        """
        if pulse_params is None:
            pulse_params = self.pulse_params
        else:
            # ensure batch dimension exists (batch-first convention)
            if len(pulse_params.shape) == 2:
                pulse_params = jnp.expand_dims(pulse_params, axis=0)
            # See note in _params_validation: never stash JAX tracers on
            # ``self``.
            if stash and not isinstance(pulse_params, jax.core.Tracer):
                self.pulse_params = pulse_params
            elif stash:
                log.debug(
                    "`pulse_params` is a JAX tracer; `self.pulse_params` is "
                    "left at its previous value."
                )

        return pulse_params

    def _enc_pulse_params_validation(
        self, enc_pulse_params: Optional[jnp.ndarray], stash: bool = True
    ) -> jnp.ndarray:
        """
        Validate and normalize encoding pulse parameters.

        Ensures encoding pulse parameters are set (using model defaults if not
        provided) and carry a leading batch dimension.

        Args:
            enc_pulse_params (Optional[jnp.ndarray]): Encoding pulse-parameter
                scalers. If None, returns the model's current encoding pulse
                parameters.

        Returns:
            jnp.ndarray: Validated encoding pulse parameters with shape
                (batch_size, n_layers, n_qubits, n_enc_pulse_per_qubit).

        Raises:
            ValueError: If the trailing dimensions do not match the model's
                encoding pulse parameter shape.
        """
        if enc_pulse_params is None:
            enc_pulse_params = self.enc_pulse_params
        else:
            # ensure batch dimension exists (batch-first convention)
            if len(enc_pulse_params.shape) == 3:
                enc_pulse_params = jnp.expand_dims(enc_pulse_params, axis=0)
            if enc_pulse_params.shape[1:] != self._enc_pulse_shape:
                raise ValueError(
                    f"enc_pulse_params trailing shape {enc_pulse_params.shape[1:]} "
                    f"does not match expected {self._enc_pulse_shape}."
                )
            # See note in _params_validation: never stash JAX tracers on
            # ``self``.
            if stash and not isinstance(enc_pulse_params, jax.core.Tracer):
                self.enc_pulse_params = enc_pulse_params

        return enc_pulse_params

    def _resolve_gate_mode(
        self,
        gate_mode: Optional[str],
        pulse_params: Optional[jnp.ndarray],
        enc_pulse_params: Optional[jnp.ndarray],
        stacklevel: int = 3,
    ) -> str:
        """
        Determine which gate groups run at pulse level.

        The mode follows from the pulse parameters that were provided:
        ``pulse_params`` lowers the ansatz and state-preparation gates,
        ``enc_pulse_params`` lowers the encoding gates, both together lower
        everything.

        Args:
            gate_mode (Optional[str]): Deprecated explicit mode. If None, the
                mode is inferred from the pulse parameters.
            pulse_params (Optional[jnp.ndarray]): Ansatz pulse-parameter scalers.
            enc_pulse_params (Optional[jnp.ndarray]): Encoding pulse-parameter
                scalers.
            stacklevel (int): Frames between this method and the user call, so
                that the deprecation warning points at the call site. Defaults
                to 3, which is correct for a direct caller.

        Returns:
            str: One of the keys of ``GATE_MODES``.

        Raises:
            ValueError: If the encoding gates would run at pulse level but the
                encoding has no pulse parametrization, or if an explicitly
                passed (deprecated) gate_mode is unknown or inconsistent with
                the provided pulse parameters.
        """
        if gate_mode is None:
            if pulse_params is not None and enc_pulse_params is not None:
                gate_mode = "all_pulse"
            elif pulse_params is not None:
                gate_mode = "ansatz_pulse"
            elif enc_pulse_params is not None:
                gate_mode = "enc_pulse"
            else:
                gate_mode = "unitary"
        else:
            warnings.warn(
                "gate_mode is deprecated, the mode is inferred from the "
                "provided pulse parameters instead. Pass pulse_params to run "
                "the ansatz at pulse level and enc_pulse_params to run the "
                "encoding at pulse level, e.g. "
                "model(pulse_params=model.pulse_params).",
                DeprecationWarning,
                stacklevel=stacklevel,
            )
            # consistency checks, only reachable via the deprecated argument
            if gate_mode not in GATE_MODES:
                raise ValueError(
                    f"Unknown gate_mode: {gate_mode}. Use one of {list(GATE_MODES)}."
                )
            if pulse_params is not None and gate_mode not in _ANSATZ_PULSE_MODES:
                raise ValueError(
                    f"pulse_params were provided but gate_mode is not one of "
                    f"{list(_ANSATZ_PULSE_MODES)}. Either switch gate_mode or do "
                    "not pass pulse_params."
                )
            if enc_pulse_params is not None and gate_mode not in _ENC_PULSE_MODES:
                raise ValueError(
                    f"enc_pulse_params were provided but gate_mode is not one of "
                    f"{list(_ENC_PULSE_MODES)}. Either switch gate_mode or do not "
                    "pass enc_pulse_params."
                )

        if gate_mode in _ENC_PULSE_MODES and not self._enc_pulse_capable:
            raise ValueError(
                "Pulse-level encoding requires an encoding whose gates have a "
                "pulse parametrization (golomb and custom callables do not). "
                "Do not pass enc_pulse_params for this model."
            )

        return gate_mode

    def _enc_params_validation(
        self, enc_params: Optional[jnp.ndarray], stash: bool = True
    ) -> jnp.ndarray:
        """
        Validate and normalize encoding parameters.

        Ensures encoding parameters have the correct shape for the model's
        input feature dimensions.

        Args:
            enc_params (Optional[jnp.ndarray]): Encoding parameters to validate.
                If None, returns the model's current encoding parameters.
            stash (bool): Whether to store the validated parameters on the
                model. See :meth:`_params_validation`. Defaults to True.

        Returns:
            jnp.ndarray: Validated encoding parameters with shape
                (n_qubits, n_input_feat).

        Raises:
            ValueError: If enc_params shape is incompatible with n_input_feat > 1.
        """
        if enc_params is None:
            enc_params = self.enc_params
        else:
            # See note in _params_validation: never stash JAX tracers on
            # ``self``.
            if stash and not isinstance(enc_params, jax.core.Tracer):
                if self.trainable_frequencies:
                    self.enc_params = enc_params
                else:
                    self.enc_params = jnp.array(enc_params)
            elif stash:
                log.debug(
                    "`enc_params` is a JAX tracer; `self.enc_params` is left "
                    "at its previous value."
                )

        if len(enc_params.shape) == 1 and self.n_input_feat == 1:
            enc_params = enc_params.reshape(-1, 1)
        elif len(enc_params.shape) == 1 and self.n_input_feat > 1:
            raise ValueError(
                f"Input dimension {self.n_input_feat} >1 but \
                `enc_params` has shape {enc_params.shape}"
            )

        return enc_params

    def _inputs_validation(
        self, inputs: Union[None, List, float, int, jnp.ndarray]
    ) -> jnp.ndarray:
        """
        Validate and normalize input data.

        Converts various input formats to a standardized 2D array shape
        suitable for batch processing in the quantum circuit.

        Args:
            inputs (Union[None, List, float, int, jnp.ndarray]): Input data in
                various formats:
                - None: Returns zeros with shape (1, n_input_feat)
                - float/int: Single scalar value
                - List: List of values or batched inputs
                - jnp.ndarray: NumPy/JAX array

        Returns:
            jnp.ndarray: Validated inputs with shape (batch_size, n_input_feat).

        Raises:
            ValueError: If input shape is incompatible with expected n_input_feat.

        Warns:
            UserWarning: If input is replicated to match n_input_feat.
        """
        if isinstance(inputs, List):
            inputs = jnp.array(np.stack(inputs))
        elif isinstance(inputs, float) or isinstance(inputs, int):
            inputs = jnp.array([inputs])
        elif inputs is None:
            inputs = jnp.array([[0] * self.n_input_feat])

        if len(inputs.shape) <= 1:
            if self.n_input_feat == 1:
                # add a batch dimension
                inputs = inputs.reshape(-1, 1)
            else:
                if inputs.shape[0] == self.n_input_feat:
                    inputs = inputs.reshape(1, -1)
                else:
                    inputs = inputs.reshape(-1, 1)
                    inputs = inputs.repeat(self.n_input_feat, axis=1)
                    warnings.warn(
                        f"Expected {self.n_input_feat} inputs, but {inputs.shape[0]} "
                        "was provided, replicating input for all input features.",
                        UserWarning,
                    )
        else:
            if inputs.shape[1] != self.n_input_feat:
                raise ValueError(
                    f"Wrong number of inputs provided. Expected {self.n_input_feat} "
                    f"inputs, but input has shape {inputs.shape}."
                )

        return inputs

    def _postprocess_res(self, result: Union[List, jnp.ndarray]) -> jnp.ndarray:
        """
        Post-process circuit execution results for uniform shape.

        Converts list outputs (from multiple measurements) to stacked arrays
        and reorders axes for consistent batch dimension placement.

        Args:
            result (Union[List, jnp.ndarray]): Raw circuit output, either a
                list of measurement results or a single array.

        Returns:
            jnp.ndarray: Uniformly shaped result array with batch dimension first.
        """
        if isinstance(result, list):
            # we use moveaxis here because in case of parity measure,
            # there is another dimension appended to the end and
            # simply transposing would result in a wrong shape
            result = jnp.stack(result)
            if len(result.shape) > 1:
                result = jnp.moveaxis(result, 0, 1)
        return result

    def _assimilate_batch(
        self,
        inputs: jnp.ndarray,
        params: jnp.ndarray,
        pulse_params: jnp.ndarray,
        enc_pulse_params: jnp.ndarray,
    ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, Tuple[int, ...]]:
        """
        Align batch dimensions across inputs, parameters, pulse parameters and
        encoding pulse parameters.

        Broadcasts and reshapes arrays to have compatible batch dimensions
        for vectorized circuit execution.

        The batch layout is ``[B_I, B_P, B_R, B_E, <payload>]`` where each array
        "owns" one batch axis and is replicated across the others (subject to
        the ``repeat_batch_axis`` mask) before being flattened to ``B``.

        Args:
            inputs (jnp.ndarray): Input data of shape (B_I, n_input_feat).
            params (jnp.ndarray): Parameters of shape (B_P, n_layers, n_params).
            pulse_params (jnp.ndarray): Pulse params of shape (B_R, n_layers, n_pulse).
            enc_pulse_params (jnp.ndarray): Encoding pulse params of shape
                (B_E, n_layers, n_qubits, n_enc_pulse).

        Returns:
            Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray,
            Tuple[int, ...]]: The four arrays, each reshaped to leading
            dimension B = B_I * B_P * B_R * B_E (subject to
            repeat_batch_axis), followed by the batch shape
            (B_I, B_P, B_R, B_E).

        Note:
            The effective batch shape depends on repeat_batch_axis configuration.
            This is the only method that derives the batch shape.
        """
        B_I = inputs.shape[0]
        # we check for the product because there is a chance that
        # there are no params. In this case we want B_P to be 1
        B_P = 1 if 0 in params.shape else params.shape[0]
        B_R = pulse_params.shape[0]
        B_E = enc_pulse_params.shape[0]

        # THIS is the only place where we derive the batch shape
        batch_shape = (B_I, B_P, B_R, B_E)
        B = np.prod(self._eff_batch_shape_of(batch_shape))

        # [B_I, ...] -> [B_I, B_P, B_R, B_E, ...] -> [B, ...]
        if B_I > 1 and self.repeat_batch_axis[0]:
            inputs = inputs[:, None, None, None, ...]
            if self.repeat_batch_axis[1]:
                inputs = jnp.repeat(inputs, B_P, axis=1)
            if self.repeat_batch_axis[2]:
                inputs = jnp.repeat(inputs, B_R, axis=2)
            if self.repeat_batch_axis[3]:
                inputs = jnp.repeat(inputs, B_E, axis=3)
            inputs = inputs.reshape(B, *inputs.shape[4:])

        # [B_P, ...] -> [B_I, B_P, B_R, B_E, ...] -> [B, ...]
        if B_P > 1 and self.repeat_batch_axis[1]:
            params = params[None, :, None, None, ...]  # [1, B_P, 1, 1, ...]
            if self.repeat_batch_axis[0]:
                params = jnp.repeat(params, B_I, axis=0)
            if self.repeat_batch_axis[2]:
                params = jnp.repeat(params, B_R, axis=2)
            if self.repeat_batch_axis[3]:
                params = jnp.repeat(params, B_E, axis=3)
            params = params.reshape(B, *params.shape[4:])

        # [B_R, ...] -> [B_I, B_P, B_R, B_E, ...] -> [B, ...]
        if B_R > 1 and self.repeat_batch_axis[2]:
            pulse_params = pulse_params[None, None, :, None, ...]  # [1, 1, B_R, 1, ...]
            if self.repeat_batch_axis[0]:
                pulse_params = jnp.repeat(pulse_params, B_I, axis=0)
            if self.repeat_batch_axis[1]:
                pulse_params = jnp.repeat(pulse_params, B_P, axis=1)
            if self.repeat_batch_axis[3]:
                pulse_params = jnp.repeat(pulse_params, B_E, axis=3)
            pulse_params = pulse_params.reshape(B, *pulse_params.shape[4:])

        # [B_E, ...] -> [B_I, B_P, B_R, B_E, ...] -> [B, ...]
        if B_E > 1 and self.repeat_batch_axis[3]:
            enc_pulse_params = enc_pulse_params[
                None, None, None, ...
            ]  # [1,1,1,B_E,...]
            if self.repeat_batch_axis[0]:
                enc_pulse_params = jnp.repeat(enc_pulse_params, B_I, axis=0)
            if self.repeat_batch_axis[1]:
                enc_pulse_params = jnp.repeat(enc_pulse_params, B_P, axis=1)
            if self.repeat_batch_axis[2]:
                enc_pulse_params = jnp.repeat(enc_pulse_params, B_R, axis=2)
            enc_pulse_params = enc_pulse_params.reshape(B, *enc_pulse_params.shape[4:])

        return inputs, params, pulse_params, enc_pulse_params, batch_shape

    def _requires_density(self) -> bool:
        """
        Check if density matrix simulation is required.

        Determines whether the circuit must be executed with the mixed-state
        simulator based on execution type and noise configuration.

        Returns:
            bool: True if density matrix simulation is required, False otherwise.
                Returns True if:
                - execution_type is "density", or
                - Any non-coherent noise channel has non-zero probability
        """
        if self.execution_type == "density":
            return True

        if self.noise_params is None:
            return False

        coherent_noise = {"GateError"}
        for k, v in self.noise_params.items():
            if k in coherent_noise:
                continue
            if v is not None and v > 0:
                return True
        return False

    def _is_stochastic(self) -> bool:
        """
        Check if execution draws random numbers at runtime.

        Only coherent gate errors and shot sampling are stochastic; the Kraus
        channels are deterministic maps on the density matrix.

        Returns:
            bool: True if the result depends on the random key.
        """
        gate_error = (self.noise_params or {}).get("GateError") or 0
        return gate_error > 0 or self.shots is not None

    @staticmethod
    def _args_are_traced(*args: Any) -> bool:
        """
        Check if any argument is a JAX tracer.

        Args:
            *args (Any): Values to inspect, may be pytrees.

        Returns:
            bool: True if the call runs inside a JAX transform.
        """
        return any(
            isinstance(x, jax.core.Tracer) for x in jax.tree_util.tree_leaves(args)
        )

    @staticmethod
    def _observable_id(obs: op.Operation) -> Any:
        """
        Get a stable identity for an observable.

        Uses the Pauli label where available and otherwise a hash of the
        matrix, memoized on the instance because reading the bytes copies the
        full $2^n \\times 2^n$ array. Mutating a matrix in place is not
        detected.

        Args:
            obs (op.Operation): Observable to identify.

        Returns:
            Any: Hashable identity of the observable.
        """
        label = getattr(obs, "_pauli_label", None)
        if label is not None:
            return label
        if getattr(obs, "_fingerprint_hash", None) is None:
            obs._fingerprint_hash = hash(np.asarray(obs.matrix).tobytes())
        return obs._fingerprint_hash

    def _structural_fingerprint(self) -> Tuple:
        """
        Summarize the circuit structure for the execution plan cache.

        Covers everything that :meth:`_variational` and :meth:`_iec` read from
        the model while recording the tape and that can change after
        initialization without changing the shapes of the execution arguments.
        Without it, a batched call would silently reuse a plan that was
        compiled for the previous structure.

        Attributes that are fixed at initialization (the encoding, the state
        preparation, the number of qubits and layers) are omitted, as replacing
        them afterwards is not supported.

        Returns:
            Tuple: Hashable structure summary, passed to
                :meth:`~jaqsi.Script.execute`.
        """
        if self._observables is None:
            obs_fingerprint = None
        else:
            obs_fingerprint = tuple(
                (o.name, tuple(o.wires), self._observable_id(o))
                for o in self._observables
            )

        return (
            self._data_reupload.shape,
            # covers the derived degree, frequencies and has_dru as well
            self._data_reupload.tobytes(),
            make_hashable(self._measured_wires),
            obs_fingerprint,
            # hashed by identity, which also covers a replaced ansatz callable
            self.pqc,
        )

    def __call__(
        self,
        params: Optional[jnp.ndarray] = None,
        inputs: Optional[jnp.ndarray] = None,
        pulse_params: Optional[jnp.ndarray] = None,
        enc_params: Optional[jnp.ndarray] = None,
        data_reupload: Union[
            bool, List[List[bool]], List[List[List[bool]]], np.ndarray
        ] = None,
        noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
        execution_type: Optional[str] = None,
        force_mean: bool = False,
        gate_mode: Optional[str] = None,
        enc_pulse_params: Optional[jnp.ndarray] = None,
        random_key: Optional[random.PRNGKey] = None,
        keepdims: bool = False,
    ) -> jnp.ndarray:
        """
        Execute the quantum circuit (callable interface).

        Provides a convenient callable interface for circuit execution,
        delegating to the _forward method.

        This method writes the arguments it receives onto the model, so it
        cannot be wrapped in an outer ``jax.jit`` or ``jax.vmap``. Use
        :meth:`apply` for that.

        Args:
            params (Optional[jnp.ndarray]): Variational parameters of shape
                (n_layers, n_params_per_layer) or (batch, n_layers, n_params_per_layer).
                If None, uses model's internal parameters.
            inputs (Optional[jnp.ndarray]): Input data of shape
                (batch_size, n_input_feat). If None, uses zero inputs.
            pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
                the ansatz and state-preparation gates. Passing them runs those
                gates at pulse level. If None, they stay unitary.
            enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
                (n_qubits, n_input_feat). If None, uses model's encoding parameters.
            data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]],
                np.ndarray]):
                Data reupload configuration. If None, uses previously set reupload
                configuration.
            noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
                Noise configuration. If None, uses previously set noise parameters.
            execution_type (Optional[str]): Measurement type: "expval", "density",
                "probs", or "state". If None, uses current execution_type setting.
            force_mean (bool): If True, averages results over measurement qubits.
                Defaults to False.
            gate_mode (Optional[str]): Deprecated. If None (default), the gate
                execution backend is inferred from the provided pulse
                parameters: ``pulse_params`` runs the ansatz and state
                preparation at pulse level, ``enc_pulse_params`` the encoding
                gates, both together everything. Passing "unitary",
                "ansatz_pulse", "enc_pulse" or "all_pulse" explicitly still
                works but emits a DeprecationWarning.
            enc_pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers
                for the encoding gates. Passing them runs the encoding gates at
                pulse level. If None, they stay unitary.
            random_key (Optional[random.PRNGKey]): JAX random key for stochastic
                execution (``GateError`` noise and shot sampling). If provided,
                the caller owns key advancement and the model's internal
                ``random_key`` is left untouched - this is the jit-safe way to
                get fresh randomness per call, since the internal key cannot be
                advanced from inside a trace. Use :meth:`next_key` to obtain
                one. If None, the internal key is used and advanced (eager
                calls only).
            keepdims (bool): If True, the full
                (B_I, B_P, B_R, B_E, O) shape is returned. If False (default),
                all singleton axes are squeezed out.

        Returns:
            jnp.ndarray: Circuit output with shape depending on execution_type:
                - "expval": (n_measured_wiress,) or scalar
                - "density": (2^n_output, 2^n_output)
                - "probs": (2^n_output,) or (n_pairs, 2^pair_size)
                - "state": (2^n_qubits,)

        Note:
            An eager call stores ``params``, ``pulse_params`` and ``enc_params``
            on the model, but a traced call (``jit``, ``grad``, ``vmap``) does
            not: JAX tracers must not outlive their transform, so the model
            state keeps its previous value. Two consequences:

            - Anything reading model state after a traced call - ``draw``,
              :class:`~qml_essentials.entanglement.Entanglement`,
              :class:`~qml_essentials.expressibility.Expressibility`, or a
              later call that omits ``params`` - sees the *old* parameters.
              Assign ``model.params = params`` yourself if the state should
              follow a traced optimization step.
            - Omitting ``params`` in a second call inside the same trace falls
              back to that stale state, so the result does not depend on the
              traced parameters (its gradient is zero). Pass ``params``
              explicitly on every call inside a trace.

            The skipped writes are reported at debug log level.
        """
        # Call forward method which handles the actual caching etc.
        return self._forward(
            params=params,
            inputs=inputs,
            pulse_params=pulse_params,
            enc_params=enc_params,
            data_reupload=data_reupload,
            noise_params=noise_params,
            execution_type=execution_type,
            force_mean=force_mean,
            gate_mode=gate_mode,
            enc_pulse_params=enc_pulse_params,
            random_key=random_key,
            keepdims=keepdims,
        )

    def apply(
        self,
        params: Optional[jnp.ndarray] = None,
        inputs: Optional[jnp.ndarray] = None,
        pulse_params: Optional[jnp.ndarray] = None,
        enc_params: Optional[jnp.ndarray] = None,
        noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
        execution_type: Optional[str] = None,
        force_mean: bool = False,
        gate_mode: Optional[str] = None,
        enc_pulse_params: Optional[jnp.ndarray] = None,
        key: Optional[random.PRNGKey] = None,
    ) -> jnp.ndarray:
        """
        Execute the quantum circuit without modifying the model.

        Functional counterpart of :meth:`__call__`. No model state is written,
        so the call can be wrapped in an outer ``jax.jit``, ``jax.vmap`` or a
        whole jitted training step. The output always keeps the full
        (B_I, B_P, B_R, B_E, O) shape, so its rank does not depend on the batch
        sizes; call ``.squeeze()`` for the shape :meth:`__call__` returns.

        Arguments left as None fall back to the current model state, which an
        outer ``jax.jit`` bakes in at trace time. Anything that varies between
        calls, such as the parameters during training or the key for shots,
        has to be passed explicitly.

        Unlike :meth:`__call__` this method takes no ``data_reupload``
        argument, as that reconfigures the circuit; set
        :attr:`data_reupload` on the model beforehand instead.

        Args:
            params (Optional[jnp.ndarray]): Variational parameters of shape
                (n_layers, n_params_per_layer) or
                (batch, n_layers, n_params_per_layer).
                If None, uses model's internal parameters.
            inputs (Optional[jnp.ndarray]): Input data of shape
                (batch_size, n_input_feat). If None, uses zero inputs.
            pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
                the ansatz and state-preparation gates. Passing them runs those
                gates at pulse level. If None, they stay unitary.
            enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
                (n_qubits, n_input_feat). If None, uses model's encoding parameters.
            noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
                Noise configuration. If None, uses the model's noise parameters.
            execution_type (Optional[str]): Measurement type: "expval", "density",
                "probs", or "state". If None, uses current execution_type setting.
            force_mean (bool): If True, averages results over measurement qubits.
                Defaults to False.
            gate_mode (Optional[str]): Deprecated. If None (default), the mode
                is inferred from the provided pulse parameters. See
                :meth:`__call__`.
            enc_pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers
                for the encoding gates. Passing them runs the encoding gates at
                pulse level. If None, they stay unitary.
            key (Optional[random.PRNGKey]): JAX random key for shots and
                stochastic noise. If None, the model's random key is used
                without advancing it.

        Returns:
            jnp.ndarray: Circuit output of shape (B_I, B_P, B_R, B_E, O), where
                O is the per-sample output shape of the execution type and may
                span more than one axis (e.g. "density" and "probs").

        Raises:
            ValueError: If the encoding gates would run at pulse level but the
                encoding has no pulse parametrization, or if shots are set for
                a density measurement.
        """
        gate_mode = self._resolve_gate_mode(gate_mode, pulse_params, enc_pulse_params)

        if execution_type is None:
            execution_type = self.execution_type
        if execution_type == "density" and self.shots is not None:
            raise ValueError("Setting execution_type to density with shots not None.")
        result_shape = self._compute_result_shape(execution_type)

        if noise_params is None:
            noise_params = self.noise_params
        else:
            noise_params = self._normalize_noise_params(noise_params)

        enc_pulse_params = self._enc_pulse_params_validation(
            enc_pulse_params, stash=False
        )
        params = self._params_validation(params, stash=False)
        pulse_params = self._pulse_params_validation(pulse_params, stash=False)
        inputs = self._inputs_validation(inputs)
        enc_params = self._enc_params_validation(enc_params, stash=False)

        inputs, params, pulse_params, enc_pulse_params, batch_shape = (
            self._assimilate_batch(
                inputs,
                params,
                pulse_params,
                enc_pulse_params,
            )
        )

        # derive a sub key as in _forward, but without advancing the model's key
        _, sub_key = safe_random_split(key if key is not None else self.random_key)

        return self._execute_forward(
            params=params,
            inputs=inputs,
            pulse_params=pulse_params,
            enc_params=enc_params,
            batch_shape=batch_shape,
            execution_type=execution_type,
            result_shape=result_shape,
            noise_params=noise_params,
            gate_mode=gate_mode,
            force_mean=force_mean,
            key=sub_key,
            enc_pulse_params=enc_pulse_params,
            keepdims=True,
        )

    def _forward(
        self,
        params: Optional[jnp.ndarray] = None,
        inputs: Optional[jnp.ndarray] = None,
        pulse_params: Optional[jnp.ndarray] = None,
        enc_params: Optional[jnp.ndarray] = None,
        data_reupload: Union[
            bool, List[List[bool]], List[List[List[bool]]], np.ndarray
        ] = None,
        noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
        execution_type: Optional[str] = None,
        force_mean: bool = False,
        gate_mode: Optional[str] = None,
        enc_pulse_params: Optional[jnp.ndarray] = None,
        random_key: Optional[random.PRNGKey] = None,
        keepdims: bool = False,
    ) -> jnp.ndarray:
        """
        Execute the quantum circuit forward pass.

        Internal implementation of the forward pass that handles parameter
        validation, batch alignment, and circuit execution routing.

        Args:
            params (Optional[jnp.ndarray]): Variational parameters of shape
                (n_layers, n_params_per_layer) or
                (batch, n_layers, n_params_per_layer).
                If None, uses model's internal parameters.
            inputs (Optional[jnp.ndarray]): Input data of shape
                (batch_size, n_input_feat).
                If None, uses zero inputs.
            pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
                pulse-mode gate execution.
            enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
                (n_qubits, n_input_feat). If None, uses model's encoding parameters.
            data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]],
                np.ndarray]):
                Data reupload configuration. If None, uses previously set reupload
                configuration.
            noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
                Noise configuration. If None, uses previously set noise parameters.
            execution_type (Optional[str]): Measurement type: "expval", "density",
                "probs", or "state". If None, uses current execution_type setting.
            force_mean (bool): If True, averages results over measurement qubits.
                Defaults to False.
            gate_mode (Optional[str]): Deprecated. If None (default), the mode
                is inferred from the provided pulse parameters. Passing
                "unitary", "ansatz_pulse", "enc_pulse" or "all_pulse"
                explicitly emits a DeprecationWarning.
            enc_pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers
                for the encoding gates. Passing them runs the encoding gates at
                pulse level. If None, they stay unitary.
            random_key (Optional[random.PRNGKey]): JAX random key for stochastic
                execution. If provided, it is used instead of (and does not
                modify) the model's internal ``random_key``. See
                :meth:`__call__` for details.
            keepdims (bool): If True, the full
                (B_I, B_P, B_R, B_E, O) shape is returned. If False (default),
                all singleton axes are squeezed out.

        Returns:
            jnp.ndarray: Circuit output with shape depending on execution_type:
                - "expval": (n_measured_wiress,) or scalar
                - "density": (2^n_output, 2^n_output)
                - "probs": (2^n_output,) or (n_pairs, 2^pair_size)
                - "state": (2^n_qubits,)

        Raises:
            ValueError: If the encoding gates would run at pulse level but the
                encoding has no pulse parametrization, or if an explicitly
                passed (deprecated) gate_mode is unknown or inconsistent with
                the provided pulse parameters.
        """
        # set the parameters as object attributes
        if noise_params is not None:
            self.noise_params = noise_params
        if execution_type is not None:
            self.execution_type = execution_type

        # one frame deeper than apply, as __call__ sits in between
        gate_mode = self._resolve_gate_mode(
            gate_mode, pulse_params, enc_pulse_params, stacklevel=4
        )

        # TODO: add testing
        if data_reupload is not None:
            self.data_reupload = data_reupload

        params = self._params_validation(params)
        pulse_params = self._pulse_params_validation(pulse_params)
        inputs = self._inputs_validation(inputs)
        enc_params = self._enc_params_validation(enc_params)
        enc_pulse_params = self._enc_pulse_params_validation(enc_pulse_params)

        inputs, params, pulse_params, enc_pulse_params, batch_shape = (
            self._assimilate_batch(
                inputs,
                params,
                pulse_params,
                enc_pulse_params,
            )
        )
        self._batch_shape = batch_shape

        # split to generate a sub_key, required for actual execution.
        if random_key is not None:
            # explicit key: purely functional, the caller advances it. This is
            # the only way to get fresh randomness inside a trace, because a
            # jitted call is traced once and then replays the trace-time key.
            _, sub_key = safe_random_split(random_key)
        else:
            if self._is_stochastic() and self._args_are_traced(
                params, inputs, pulse_params, enc_pulse_params
            ):
                warnings.warn(
                    "Stochastic execution (`GateError` or `shots`) without an "
                    "explicit `random_key` inside a JAX transform: the key is "
                    "read at trace time, so a jitted function replays the same "
                    "noise realization on every call. Pass "
                    "`random_key=model.next_key()` from outside the transform.",
                    UserWarning,
                )
            # Under JAX tracing (jit) the split result is a tracer; stashing it
            # on ``self`` leaks the tracer across calls (UnexpectedTracerError),
            # so only advance the key eagerly. Note that a jitted call
            # therefore reuses the same key on every execution - pass
            # ``random_key`` explicitly if that matters.
            new_key, sub_key = safe_random_split(self.random_key)
            if not isinstance(new_key, jax.core.Tracer):
                self.random_key = new_key

        return self._execute_forward(
            params=params,
            inputs=inputs,
            pulse_params=pulse_params,
            enc_params=enc_params,
            batch_shape=batch_shape,
            enc_pulse_params=enc_pulse_params,
            execution_type=self.execution_type,
            result_shape=self._result_shape,
            noise_params=self.noise_params,
            gate_mode=gate_mode,
            force_mean=force_mean,
            key=sub_key,
            keepdims=keepdims,
        )

    def _execute_forward(
        self,
        params: jnp.ndarray,
        inputs: jnp.ndarray,
        pulse_params: jnp.ndarray,
        enc_params: jnp.ndarray,
        enc_pulse_params: jnp.ndarray,
        batch_shape: Tuple[int, ...],
        execution_type: str,
        result_shape: Tuple[int, ...],
        noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]],
        gate_mode: str,
        force_mean: bool,
        key: random.PRNGKey,
        keepdims: bool,
    ) -> jnp.ndarray:
        """
        Run the simulation and post-process the result.

        This is the shared core of :meth:`_forward` and :meth:`apply`. Every
        value it needs is passed in explicitly and no model state is written,
        so it is safe to call from within an outer JAX transform.

        Args:
            params (jnp.ndarray): Validated and batch-aligned parameters.
            inputs (jnp.ndarray): Validated and batch-aligned inputs.
            pulse_params (jnp.ndarray): Validated and batch-aligned pulse params.
            enc_params (jnp.ndarray): Validated encoding parameters.
            enc_pulse_params (jnp.ndarray): Validated and batch-aligned
                encoding pulse parameters.
            batch_shape (Tuple[int, ...]): Batch shape (B_I, B_P, B_R, B_E).
            execution_type (str): Measurement type: "expval", "density",
                "probs", or "state".
            result_shape (Tuple[int, ...]): Per-sample output shape.
            noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
                Normalized noise configuration.
            gate_mode (str): Gate execution backend, "unitary",
                "ansatz_pulse", "enc_pulse" or "all_pulse".
            force_mean (bool): If True, averages results over the output axis.
            key (random.PRNGKey): JAX random key for this execution.
            keepdims (bool): If True, the full (B_I, B_P, B_R, B_E, O) shape is
                returned. If False, all singleton axes are squeezed out.

        Returns:
            jnp.ndarray: Circuit output.
        """
        # Build measurement type & observables from execution_type / output_qubit
        meas_type, obs = self._build_obs(execution_type)

        eff_batch_shape = self._eff_batch_shape_of(batch_shape)

        # Jaqsi auto-routes between statevector and density-matrix simulation
        # based on whether noise channels appear on the tape, so a single
        B = np.prod(eff_batch_shape)

        # kwargs are broadcast (not vmapped over)
        exec_kwargs = dict(
            noise_params=noise_params,
            gate_mode=gate_mode,
        )

        # Build a shot key from the given key if shots are requested
        shot_key = None
        sub_key = key
        if self.shots is not None:
            # overwrite subkey and split shot_key
            sub_key, shot_key = safe_random_split(sub_key)

        if B > 1:
            # use random keys, derived from the subkey
            random_keys = safe_random_split(sub_key, num=B)

            in_axes = (
                0 if batch_shape[1] > 1 else None,  # params
                0 if batch_shape[0] > 1 else None,  # inputs
                0 if batch_shape[2] > 1 else None,  # pulse_params
                0,  # random_keys
                None,  # enc_params (broadcast, not batched)
                0 if batch_shape[3] > 1 else None,  # enc_pulse_params
            )

            result = self.script.execute(
                type=meas_type,
                obs=obs,
                args=(
                    params,
                    inputs,
                    pulse_params,
                    random_keys,
                    enc_params,
                    enc_pulse_params,
                ),
                kwargs=exec_kwargs,
                in_axes=in_axes,
                shots=self.shots,
                key=shot_key,
                fingerprint=self._structural_fingerprint(),
            )
        else:
            # use the subkey directly
            result = self.script.execute(
                type=meas_type,
                obs=obs,
                args=(
                    params,
                    inputs,
                    pulse_params,
                    sub_key,
                    enc_params,
                    enc_pulse_params,
                ),
                kwargs=exec_kwargs,
                shots=self.shots,
                key=shot_key,
                fingerprint=self._structural_fingerprint(),
            )

        result = self._postprocess_res(result)

        # --- Post-processing for partial-qubit measurements ---------------
        if execution_type == "density" and not self.all_qubit_measurement:
            result = js.partial_trace(result, self.n_qubits, self._measured_wires)

        if execution_type == "probs" and not self.all_qubit_measurement:
            if isinstance(self._measured_wires[0], (list, tuple)):
                # list of qubit groups - marginalize each independently
                result = jnp.stack(
                    [
                        js.marginalize_probs(result, self.n_qubits, list(group))
                        for group in self._measured_wires
                    ]
                )
            else:
                result = js.marginalize_probs(
                    result, self.n_qubits, self._measured_wires
                )

        result = jnp.asarray(result)
        result = result.reshape((*eff_batch_shape, *result_shape))
        if not keepdims:
            result = result.squeeze()

        if (
            execution_type in ("expval", "probs")
            and force_mean
            and len(result.shape) > 0
            and result_shape[0] > 1
        ):
            result = result.mean(axis=-1, keepdims=keepdims)

        return result

`all_qubit_measurement` `property` #

Check if measurement is performed on all qubits.

`batch_shape` `property` #

Get the batch shape (B_I, B_P, B_R, B_E). If the model was not called before, it returns (1, 1, 1, 1).

Returns:

Type	Description
`Tuple[int, ...]`	Tuple[int, ...]: Tuple of (input_batch, param_batch, pulse_batch, enc_pulse_batch). Returns (1, 1, 1, 1) if model has not been called yet.

`data_reupload` `property` `writable` #

Get the data reupload mask.

`degree` `property` `writable` #

Get the degree of the model.

`eff_batch_shape` `property` #

Get the effective batch shape after applying repeat_batch_axis mask.

Returns:

Type	Description
`Tuple[int, ...]`	Tuple[int, ...]: Effective batch dimensions, excluding zeros.

`enc_params` `property` `writable` #

Get the encoding parameters used for input transformation.

`enc_pulse_params` `property` `writable` #

Get the encoding pulse parameters for all_pulse-mode execution.

`execution_type` `property` `writable` #

Gets the execution type of the model.

Returns:

Name	Type	Description
`str`	`str`	The execution type, one of 'density', 'expval', or 'probs'.

`frequencies` `property` `writable` #

Get the frequencies of the model.

`has_dru` `property` #

Check if the model has data reupload.

`noise_params` `property` `writable` #

Gets the noise parameters of the model.

Returns:

Type	Description
`Optional[Dict[str, Union[float, Dict[str, float]]]]`	Optional[Dict[str, float]]: A dictionary of
`Optional[Dict[str, Union[float, Dict[str, float]]]]`	noise parameters or None if not set.

`observables` `property` `writable` #

The custom :class:~jaqsi.operations.Operation observables, or the list of measured wires when using the default PauliZ readout.

With a list of observables, __call__ and execution_type="expval" returns one expectation value per observable instead of one PauliZ per measured qubit.

`output_qubit` `property` `writable` #

Deprecated alias for :attr:observables; returns the measured wires.

`params` `property` `writable` #

Get the variational parameters of the model.

`pulse_params` `property` `writable` #

Get the pulse parameters for pulse-mode gate execution.

`shots` `property` `writable` #

Gets the number of shots to use for the quantum device.

Returns:

Type	Description
`Optional[int]`	Optional[int]: The number of shots.

`call(params=None, inputs=None, pulse_params=None, enc_params=None, data_reupload=None, noise_params=None, execution_type=None, force_mean=False, gate_mode=None, enc_pulse_params=None, random_key=None, keepdims=False)` #

Execute the quantum circuit (callable interface).

Provides a convenient callable interface for circuit execution, delegating to the _forward method.

This method writes the arguments it receives onto the model, so it cannot be wrapped in an outer jax.jit or jax.vmap. Use :meth:apply for that.

Parameters:

Name	Type	Description	Default
`params`	`Optional[ndarray]`	Variational parameters of shape (n_layers, n_params_per_layer) or (batch, n_layers, n_params_per_layer). If None, uses model's internal parameters.	`None`
`inputs`	`Optional[ndarray]`	Input data of shape (batch_size, n_input_feat). If None, uses zero inputs.	`None`
`pulse_params`	`Optional[ndarray]`	Pulse parameter scalers for the ansatz and state-preparation gates. Passing them runs those gates at pulse level. If None, they stay unitary.	`None`
`enc_params`	`Optional[ndarray]`	Encoding parameters of shape (n_qubits, n_input_feat). If None, uses model's encoding parameters.	`None`
`noise_params`	`Optional[Dict[str, Union[float, Dict[str, float]]]]`	Noise configuration. If None, uses previously set noise parameters.	`None`
`execution_type`	`Optional[str]`	Measurement type: "expval", "density", "probs", or "state". If None, uses current execution_type setting.	`None`
`force_mean`	`bool`	If True, averages results over measurement qubits. Defaults to False.	`False`
`gate_mode`	`Optional[str]`	Deprecated. If None (default), the gate execution backend is inferred from the provided pulse parameters: `pulse_params` runs the ansatz and state preparation at pulse level, `enc_pulse_params` the encoding gates, both together everything. Passing "unitary", "ansatz_pulse", "enc_pulse" or "all_pulse" explicitly still works but emits a DeprecationWarning.	`None`
`enc_pulse_params`	`Optional[ndarray]`	Pulse parameter scalers for the encoding gates. Passing them runs the encoding gates at pulse level. If None, they stay unitary.	`None`
`random_key`	`Optional[PRNGKey]`	JAX random key for stochastic execution (`GateError` noise and shot sampling). If provided, the caller owns key advancement and the model's internal `random_key` is left untouched - this is the jit-safe way to get fresh randomness per call, since the internal key cannot be advanced from inside a trace. Use :meth:`next_key` to obtain one. If None, the internal key is used and advanced (eager calls only).	`None`
`keepdims`	`bool`	If True, the full (B_I, B_P, B_R, B_E, O) shape is returned. If False (default), all singleton axes are squeezed out.	`False`

Returns:

Type	Description
`ndarray`	jnp.ndarray: Circuit output with shape depending on execution_type: - "expval": (n_measured_wiress,) or scalar - "density": (2^n_output, 2^n_output) - "probs": (2^n_output,) or (n_pairs, 2^pair_size) - "state": (2^n_qubits,)

Note

An eager call stores params, pulse_params and enc_params on the model, but a traced call (jit, grad, vmap) does not: JAX tracers must not outlive their transform, so the model state keeps its previous value. Two consequences:

Anything reading model state after a traced call - draw, :class:~qml_essentials.entanglement.Entanglement, :class:~qml_essentials.expressibility.Expressibility, or a later call that omits params - sees the old parameters. Assign model.params = params yourself if the state should follow a traced optimization step.
Omitting params in a second call inside the same trace falls back to that stale state, so the result does not depend on the traced parameters (its gradient is zero). Pass params explicitly on every call inside a trace.

The skipped writes are reported at debug log level.

Source code in qml_essentials/model.py

def __call__(
    self,
    params: Optional[jnp.ndarray] = None,
    inputs: Optional[jnp.ndarray] = None,
    pulse_params: Optional[jnp.ndarray] = None,
    enc_params: Optional[jnp.ndarray] = None,
    data_reupload: Union[
        bool, List[List[bool]], List[List[List[bool]]], np.ndarray
    ] = None,
    noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
    execution_type: Optional[str] = None,
    force_mean: bool = False,
    gate_mode: Optional[str] = None,
    enc_pulse_params: Optional[jnp.ndarray] = None,
    random_key: Optional[random.PRNGKey] = None,
    keepdims: bool = False,
) -> jnp.ndarray:
    """
    Execute the quantum circuit (callable interface).

    Provides a convenient callable interface for circuit execution,
    delegating to the _forward method.

    This method writes the arguments it receives onto the model, so it
    cannot be wrapped in an outer ``jax.jit`` or ``jax.vmap``. Use
    :meth:`apply` for that.

    Args:
        params (Optional[jnp.ndarray]): Variational parameters of shape
            (n_layers, n_params_per_layer) or (batch, n_layers, n_params_per_layer).
            If None, uses model's internal parameters.
        inputs (Optional[jnp.ndarray]): Input data of shape
            (batch_size, n_input_feat). If None, uses zero inputs.
        pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
            the ansatz and state-preparation gates. Passing them runs those
            gates at pulse level. If None, they stay unitary.
        enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
            (n_qubits, n_input_feat). If None, uses model's encoding parameters.
        data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]],
            np.ndarray]):
            Data reupload configuration. If None, uses previously set reupload
            configuration.
        noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
            Noise configuration. If None, uses previously set noise parameters.
        execution_type (Optional[str]): Measurement type: "expval", "density",
            "probs", or "state". If None, uses current execution_type setting.
        force_mean (bool): If True, averages results over measurement qubits.
            Defaults to False.
        gate_mode (Optional[str]): Deprecated. If None (default), the gate
            execution backend is inferred from the provided pulse
            parameters: ``pulse_params`` runs the ansatz and state
            preparation at pulse level, ``enc_pulse_params`` the encoding
            gates, both together everything. Passing "unitary",
            "ansatz_pulse", "enc_pulse" or "all_pulse" explicitly still
            works but emits a DeprecationWarning.
        enc_pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers
            for the encoding gates. Passing them runs the encoding gates at
            pulse level. If None, they stay unitary.
        random_key (Optional[random.PRNGKey]): JAX random key for stochastic
            execution (``GateError`` noise and shot sampling). If provided,
            the caller owns key advancement and the model's internal
            ``random_key`` is left untouched - this is the jit-safe way to
            get fresh randomness per call, since the internal key cannot be
            advanced from inside a trace. Use :meth:`next_key` to obtain
            one. If None, the internal key is used and advanced (eager
            calls only).
        keepdims (bool): If True, the full
            (B_I, B_P, B_R, B_E, O) shape is returned. If False (default),
            all singleton axes are squeezed out.

    Returns:
        jnp.ndarray: Circuit output with shape depending on execution_type:
            - "expval": (n_measured_wiress,) or scalar
            - "density": (2^n_output, 2^n_output)
            - "probs": (2^n_output,) or (n_pairs, 2^pair_size)
            - "state": (2^n_qubits,)

    Note:
        An eager call stores ``params``, ``pulse_params`` and ``enc_params``
        on the model, but a traced call (``jit``, ``grad``, ``vmap``) does
        not: JAX tracers must not outlive their transform, so the model
        state keeps its previous value. Two consequences:

        - Anything reading model state after a traced call - ``draw``,
          :class:`~qml_essentials.entanglement.Entanglement`,
          :class:`~qml_essentials.expressibility.Expressibility`, or a
          later call that omits ``params`` - sees the *old* parameters.
          Assign ``model.params = params`` yourself if the state should
          follow a traced optimization step.
        - Omitting ``params`` in a second call inside the same trace falls
          back to that stale state, so the result does not depend on the
          traced parameters (its gradient is zero). Pass ``params``
          explicitly on every call inside a trace.

        The skipped writes are reported at debug log level.
    """
    # Call forward method which handles the actual caching etc.
    return self._forward(
        params=params,
        inputs=inputs,
        pulse_params=pulse_params,
        enc_params=enc_params,
        data_reupload=data_reupload,
        noise_params=noise_params,
        execution_type=execution_type,
        force_mean=force_mean,
        gate_mode=gate_mode,
        enc_pulse_params=enc_pulse_params,
        random_key=random_key,
        keepdims=keepdims,
    )

`init(n_qubits, n_layers, circuit_type='No_Ansatz', data_reupload=True, state_preparation=None, encoding=Gates.RX, trainable_frequencies=False, initialization='random', initialization_domain=[0, 2 * jnp.pi], output_qubit=None, observables=None, shots=None, random_seed=1000, repeat_batch_axis=[True, True, True, True], pulse_shape='gaussian')` #

Initialize the quantum circuit model. Parameters will have the shape [impl_n_layers, parameters_per_layer] where impl_n_layers is the number of layers provided and added by one depending if data_reupload is True and parameters_per_layer is given by the chosen ansatz.

The model is initialized with the following parameters as defaults: - noise_params: None - execution_type: "expval" - shots: None

Parameters:

Name	Type	Description	Default
`n_qubits`	`int`	The number of qubits in the circuit.	required
`n_layers`	`int`	The number of layers in the circuit.	required
`circuit_type`	`(str, Circuit)`	The type of quantum circuit to use. If None, defaults to "no_ansatz".	`'No_Ansatz'`
`encoding`	`Union[str, Callable, List[str], List[Callable]]`	The unitary to use for encoding the input data. Can be a string (e.g. "RX") or a callable (e.g. gateset.RX). Defaults to gateset.RX. If input is multidimensional it is assumed to be a list of unitaries or a list of strings.	`RX`
`trainable_frequencies`	`bool`	Sets trainable encoding parameters for trainable frequencies. Defaults to False.	`False`
`initialization`	`str`	The strategy to initialize the parameters. Can be "random", "zeros", "zero-controlled", "pi", or "pi-controlled". Defaults to "random".	`'random'`
`output_qubit`	`(List[int], int)`	Deprecated alias for `observables`. Forwards to `observables` and will be removed in a future release. Defaults to None.	`None`
`shots`	`Optional[int]`	The number of shots to use for the quantum device. Defaults to None.	`None`
`random_seed`	`int`	seed for the random number generator in initialization is "random" and for random noise parameters. Defaults to 1000.	`1000`
`repeat_batch_axis`	`List[bool]`	Each boolean in the array determines over which axes to parallelise computation. The axes correspond to [inputs, params, pulse_params, enc_pulse_params]. Defaults to [True, True, True, True], meaning that batching is enabled over all axes. A 3-element list (legacy) is accepted and extended with a trailing True for the enc_pulse_params axis.	`[True, True, True, True]`
`pulse_shape`	`str`	Pulse envelope shape for pulse-level simulation. One of `PulseEnvelope.available()`. Defaults to `"gaussian"`.	`'gaussian'`

Returns:

Type	Description
`None`	None

Source code in qml_essentials/model.py

def __init__(
    self,
    n_qubits: int,
    n_layers: int,
    circuit_type: Union[str, type[Circuit]] = "No_Ansatz",
    data_reupload: Union[
        bool, List[List[bool]], List[List[List[bool]]], np.ndarray
    ] = True,
    state_preparation: Union[
        str, Callable, List[Union[str, Callable]], None
    ] = None,
    encoding: Union[Encoding, str, Callable, List[Union[str, Callable]]] = Gates.RX,
    trainable_frequencies: bool = False,
    initialization: str = "random",
    initialization_domain: List[float] = [0, 2 * jnp.pi],
    output_qubit: Union[List[int], int, None] = None,
    observables: Union[
        int, List[Union[int, List[int]]], List[op.Operation], None
    ] = None,
    shots: Optional[int] = None,
    random_seed: int = 1000,
    repeat_batch_axis: List[bool] = [True, True, True, True],
    pulse_shape: str = "gaussian",
) -> None:
    """
    Initialize the quantum circuit model.
    Parameters will have the shape [impl_n_layers, parameters_per_layer]
    where impl_n_layers is the number of layers provided and added by one
    depending if data_reupload is True and parameters_per_layer is given by
    the chosen ansatz.

    The model is initialized with the following parameters as defaults:
    - noise_params: None
    - execution_type: "expval"
    - shots: None

    Args:
        n_qubits (int): The number of qubits in the circuit.
        n_layers (int): The number of layers in the circuit.
        circuit_type (str, Circuit): The type of quantum circuit to use.
            If None, defaults to "no_ansatz".
        data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]],
            np.ndarray], optional):
            Whether to reupload data to the quantum device on each
            layer and qubit. Detailed re-uploading instructions can be given
            as a list/array of 0/False and 1/True with shape (n_qubits,
            n_layers) to specify where to upload the data. Defaults to True
            for applying data re-uploading to the full circuit.
        encoding (Union[str, Callable, List[str], List[Callable]], optional):
            The unitary to use for encoding the input data. Can be a string
            (e.g. "RX") or a callable (e.g. gateset.RX). Defaults to gateset.RX.
            If input is multidimensional it is assumed to be a list of
            unitaries or a list of strings.
        trainable_frequencies (bool, optional):
            Sets trainable encoding parameters for trainable frequencies.
            Defaults to False.
        initialization (str, optional): The strategy to initialize the parameters.
            Can be "random", "zeros", "zero-controlled", "pi", or "pi-controlled".
            Defaults to "random".
        output_qubit (List[int], int, optional): Deprecated alias for
            ``observables``. Forwards to ``observables`` and will be removed
            in a future release. Defaults to None.
        observables (int, List[int], List[List[int]], List[op.Operation],
            optional): Measurement specification. A qubit index, a list of
            indices, or a list of qubit groups (for $Z$-parity) selects the
            measured subsystem with the default PauliZ readout.
            Alternatively, a list of
            :class:`~jaqsi.operations.Operation` observables makes
            ``execution_type="expval"`` return one expectation value per
            observable. When None all qubits are measured. Defaults to None.
        shots (Optional[int], optional): The number of shots to use for
            the quantum device. Defaults to None.
        random_seed (int, optional): seed for the random number generator
            in initialization is "random" and for random noise parameters.
            Defaults to 1000.
        repeat_batch_axis (List[bool], optional): Each boolean in the array
            determines over which axes to parallelise computation. The axes
            correspond to [inputs, params, pulse_params, enc_pulse_params].
            Defaults to [True, True, True, True], meaning that batching is
            enabled over all axes. A 3-element list (legacy) is accepted and
            extended with a trailing True for the enc_pulse_params axis.
        pulse_shape (str, optional): Pulse envelope shape for pulse-level
            simulation. One of ``PulseEnvelope.available()``.
            Defaults to ``"gaussian"``.

    Returns:
        None
    """
    # Initialize default parameters needed for circuit evaluation
    self.n_qubits: int = n_qubits
    if output_qubit is not None:
        if observables is not None:
            raise ValueError("Pass either output_qubit or observables, not both.")
        warnings.warn(
            "output_qubit is deprecated, use observables instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        observables = output_qubit
    self.observables = observables
    self.n_layers: int = n_layers
    self.noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None
    self.shots = shots
    self.trainable_frequencies: bool = trainable_frequencies
    self.execution_type: str = "expval"
    # backward compatibility
    # TODO: consider making this more generic in future
    # (for someone wanting to control this without bothering with pulse stuff)
    if len(repeat_batch_axis) == 3:
        log.warning("Batch axis should have length 4")
        repeat_batch_axis = list(repeat_batch_axis) + [True]
    self.repeat_batch_axis: List[bool] = repeat_batch_axis

    # --- Pulse envelope ---
    pinfo.set_envelope(pulse_shape)

    # --- State Preparation ---
    try:
        self._sp = Gates.parse_gates(state_preparation, Gates)
    except ValueError as e:
        raise ValueError(f"Error parsing encodings: {e}")

    # prepare corresponding pulse parameters (always optimized pulses)
    self.sp_pulse_params = []
    for sp in self._sp:
        sp_name = sp.__name__ if hasattr(sp, "__name__") else str(sp)

        if pinfo.gate_by_name(sp_name) is not None:
            self.sp_pulse_params.append(pinfo.gate_by_name(sp_name).params)
        else:
            # gate has no pulse parametrization
            self.sp_pulse_params.append(None)

    # --- Encoding ---
    if isinstance(encoding, Encoding):
        # user wants custom strategy? do it!
        self._enc = encoding
    else:
        # use hammming encoding by default
        self._enc = Encoding("hamming", encoding)

    if self._enc.is_golomb:
        self._enc._n_qubits = n_qubits

    # Number of possible inputs
    self.n_input_feat = len(self._enc)
    log.debug(f"Number of input features: {self.n_input_feat}")

    # Trainable frequencies, default initialization as in arXiv:2309.03279v2
    self.enc_params = jnp.ones((self.n_layers, self.n_qubits, self.n_input_feat))

    # Per-feature pulse-parameter sizes/offsets used to slice
    # enc_pulse_params in _iec under "all_pulse" mode. Only encodings whose
    # gates all have a pulse parametrization are supported (golomb and
    # custom callables do not).
    # TODO: golomb should be doable but needs a closer investigation
    self._enc_pulse_sizes: List[int] = []
    self._enc_pulse_capable = not self._enc.is_golomb
    if self._enc_pulse_capable:
        for g in self._enc._gates:
            if pinfo.gate_by_name(g) is None:
                self._enc_pulse_capable = False
                self._enc_pulse_sizes = []
                break
            self._enc_pulse_sizes.append(pinfo.gate_by_name(g).size)

    self._enc_pulse_offsets: List[int] = list(
        np.cumsum([0, *self._enc_pulse_sizes[:-1]])
    )
    self._enc_pulse_shape: Tuple[int, int, int] = (
        self.n_layers,
        self.n_qubits,
        sum(self._enc_pulse_sizes),
    )

    # --- Data-Reuploading ---

    # Keep as NumPy array (not JAX) so that ``if data_reupload[q, idx]``
    # in _iec remains a concrete Python bool even under jax.jit tracing.
    # note that setting this will also update self.degree and self.frequencies
    # and in consequence also self.has_dru
    self.data_reupload = data_reupload

    # check for the highest degree among all input dimensions
    if self.has_dru:
        impl_n_layers: int = n_layers + 1  # we need L+1 according to Schuld et al.
    else:
        impl_n_layers = n_layers
    log.info(f"Number of implicit layers: {impl_n_layers}.")

    # --- Ansatz ---
    # only weak check for str. We trust the user to provide sth useful
    if isinstance(circuit_type, str):
        self.pqc: Callable[[Optional[jnp.ndarray], int], int] = getattr(
            Ansaetze, circuit_type or "No_Ansatz"
        )()
    else:
        self.pqc = circuit_type()
    log.info(f"Using Ansatz {circuit_type}.")

    # calculate the shape of the parameter vector here, we will re-use this in init.
    params_per_layer = self.pqc.n_params_per_layer(self.n_qubits)
    self._params_shape: Tuple[int, int] = (impl_n_layers, params_per_layer)
    log.info(f"Parameters per layer: {params_per_layer}")

    pulse_params_per_layer = self.pqc.n_pulse_params_per_layer(self.n_qubits)
    self._pulse_params_shape: Tuple[int, int] = (
        impl_n_layers,
        pulse_params_per_layer,
    )

    # intialize to None as we can't know this yet
    self._batch_shape = None

    # this will also be re-used in the init method,
    # however, only if nothing is provided
    self._inialization_strategy = initialization
    self._initialization_domain = initialization_domain

    # ..here! where we only require a JAX random key
    self.random_key = self.initialize_params(random.key(random_seed))

    # Initializing pulse params
    self.pulse_params: jnp.ndarray = jnp.ones((1, *self._pulse_params_shape))

    log.info(f"Initialized pulse parameters with shape {self.pulse_params.shape}.")

    # Initializing encoding pulse params (element-wise scalers, ones by
    # default). Batch-first convention, mirroring pulse_params.
    self.enc_pulse_params: jnp.ndarray = jnp.ones((1, *self._enc_pulse_shape))

    log.info(
        f"Initialized encoding pulse parameters with shape "
        f"{self.enc_pulse_params.shape}."
    )

    # Initialise the jaqsi Script that wraps _variational.
    # No device selection needed - jaqsi auto-routes between statevector
    # and density-matrix simulation based on whether noise channels are
    # present on the tape.
    self.script = js.Script(f=self._variational, n_qubits=self.n_qubits)

`repr()` #

Return text representation of the quantum circuit model.

Source code in qml_essentials/model.py

def __repr__(self) -> str:
    """Return text representation of the quantum circuit model."""
    return self.draw(figure="text")

`str()` #

Return string representation of the quantum circuit model.

Source code in qml_essentials/model.py

def __str__(self) -> str:
    """Return string representation of the quantum circuit model."""
    return self.draw(figure="text")

`apply(params=None, inputs=None, pulse_params=None, enc_params=None, noise_params=None, execution_type=None, force_mean=False, gate_mode=None, enc_pulse_params=None, key=None)` #

Execute the quantum circuit without modifying the model.

Functional counterpart of :meth:__call__. No model state is written, so the call can be wrapped in an outer jax.jit, jax.vmap or a whole jitted training step. The output always keeps the full (B_I, B_P, B_R, B_E, O) shape, so its rank does not depend on the batch sizes; call .squeeze() for the shape :meth:__call__ returns.

Arguments left as None fall back to the current model state, which an outer jax.jit bakes in at trace time. Anything that varies between calls, such as the parameters during training or the key for shots, has to be passed explicitly.

Unlike :meth:__call__ this method takes no data_reupload argument, as that reconfigures the circuit; set :attr:data_reupload on the model beforehand instead.

Parameters:

Name	Type	Description	Default
`params`	`Optional[ndarray]`	Variational parameters of shape (n_layers, n_params_per_layer) or (batch, n_layers, n_params_per_layer). If None, uses model's internal parameters.	`None`
`inputs`	`Optional[ndarray]`	Input data of shape (batch_size, n_input_feat). If None, uses zero inputs.	`None`
`pulse_params`	`Optional[ndarray]`	Pulse parameter scalers for the ansatz and state-preparation gates. Passing them runs those gates at pulse level. If None, they stay unitary.	`None`
`enc_params`	`Optional[ndarray]`	Encoding parameters of shape (n_qubits, n_input_feat). If None, uses model's encoding parameters.	`None`
`noise_params`	`Optional[Dict[str, Union[float, Dict[str, float]]]]`	Noise configuration. If None, uses the model's noise parameters.	`None`
`execution_type`	`Optional[str]`	Measurement type: "expval", "density", "probs", or "state". If None, uses current execution_type setting.	`None`
`force_mean`	`bool`	If True, averages results over measurement qubits. Defaults to False.	`False`
`gate_mode`	`Optional[str]`	Deprecated. If None (default), the mode is inferred from the provided pulse parameters. See :meth:`__call__`.	`None`
`enc_pulse_params`	`Optional[ndarray]`	Pulse parameter scalers for the encoding gates. Passing them runs the encoding gates at pulse level. If None, they stay unitary.	`None`
`key`	`Optional[PRNGKey]`	JAX random key for shots and stochastic noise. If None, the model's random key is used without advancing it.	`None`

Returns:

Type	Description
`ndarray`	jnp.ndarray: Circuit output of shape (B_I, B_P, B_R, B_E, O), where O is the per-sample output shape of the execution type and may span more than one axis (e.g. "density" and "probs").

Raises:

Type	Description
`ValueError`	If the encoding gates would run at pulse level but the encoding has no pulse parametrization, or if shots are set for a density measurement.

Source code in qml_essentials/model.py

def apply(
    self,
    params: Optional[jnp.ndarray] = None,
    inputs: Optional[jnp.ndarray] = None,
    pulse_params: Optional[jnp.ndarray] = None,
    enc_params: Optional[jnp.ndarray] = None,
    noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
    execution_type: Optional[str] = None,
    force_mean: bool = False,
    gate_mode: Optional[str] = None,
    enc_pulse_params: Optional[jnp.ndarray] = None,
    key: Optional[random.PRNGKey] = None,
) -> jnp.ndarray:
    """
    Execute the quantum circuit without modifying the model.

    Functional counterpart of :meth:`__call__`. No model state is written,
    so the call can be wrapped in an outer ``jax.jit``, ``jax.vmap`` or a
    whole jitted training step. The output always keeps the full
    (B_I, B_P, B_R, B_E, O) shape, so its rank does not depend on the batch
    sizes; call ``.squeeze()`` for the shape :meth:`__call__` returns.

    Arguments left as None fall back to the current model state, which an
    outer ``jax.jit`` bakes in at trace time. Anything that varies between
    calls, such as the parameters during training or the key for shots,
    has to be passed explicitly.

    Unlike :meth:`__call__` this method takes no ``data_reupload``
    argument, as that reconfigures the circuit; set
    :attr:`data_reupload` on the model beforehand instead.

    Args:
        params (Optional[jnp.ndarray]): Variational parameters of shape
            (n_layers, n_params_per_layer) or
            (batch, n_layers, n_params_per_layer).
            If None, uses model's internal parameters.
        inputs (Optional[jnp.ndarray]): Input data of shape
            (batch_size, n_input_feat). If None, uses zero inputs.
        pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
            the ansatz and state-preparation gates. Passing them runs those
            gates at pulse level. If None, they stay unitary.
        enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
            (n_qubits, n_input_feat). If None, uses model's encoding parameters.
        noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
            Noise configuration. If None, uses the model's noise parameters.
        execution_type (Optional[str]): Measurement type: "expval", "density",
            "probs", or "state". If None, uses current execution_type setting.
        force_mean (bool): If True, averages results over measurement qubits.
            Defaults to False.
        gate_mode (Optional[str]): Deprecated. If None (default), the mode
            is inferred from the provided pulse parameters. See
            :meth:`__call__`.
        enc_pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers
            for the encoding gates. Passing them runs the encoding gates at
            pulse level. If None, they stay unitary.
        key (Optional[random.PRNGKey]): JAX random key for shots and
            stochastic noise. If None, the model's random key is used
            without advancing it.

    Returns:
        jnp.ndarray: Circuit output of shape (B_I, B_P, B_R, B_E, O), where
            O is the per-sample output shape of the execution type and may
            span more than one axis (e.g. "density" and "probs").

    Raises:
        ValueError: If the encoding gates would run at pulse level but the
            encoding has no pulse parametrization, or if shots are set for
            a density measurement.
    """
    gate_mode = self._resolve_gate_mode(gate_mode, pulse_params, enc_pulse_params)

    if execution_type is None:
        execution_type = self.execution_type
    if execution_type == "density" and self.shots is not None:
        raise ValueError("Setting execution_type to density with shots not None.")
    result_shape = self._compute_result_shape(execution_type)

    if noise_params is None:
        noise_params = self.noise_params
    else:
        noise_params = self._normalize_noise_params(noise_params)

    enc_pulse_params = self._enc_pulse_params_validation(
        enc_pulse_params, stash=False
    )
    params = self._params_validation(params, stash=False)
    pulse_params = self._pulse_params_validation(pulse_params, stash=False)
    inputs = self._inputs_validation(inputs)
    enc_params = self._enc_params_validation(enc_params, stash=False)

    inputs, params, pulse_params, enc_pulse_params, batch_shape = (
        self._assimilate_batch(
            inputs,
            params,
            pulse_params,
            enc_pulse_params,
        )
    )

    # derive a sub key as in _forward, but without advancing the model's key
    _, sub_key = safe_random_split(key if key is not None else self.random_key)

    return self._execute_forward(
        params=params,
        inputs=inputs,
        pulse_params=pulse_params,
        enc_params=enc_params,
        batch_shape=batch_shape,
        execution_type=execution_type,
        result_shape=result_shape,
        noise_params=noise_params,
        gate_mode=gate_mode,
        force_mean=force_mean,
        key=sub_key,
        enc_pulse_params=enc_pulse_params,
        keepdims=True,
    )

`draw(inputs=None, figure='text', **kwargs)` #

Visualize the quantum circuit.

Records the circuit tape (without noise) and renders the gate sequence using the requested backend.

Parameters:

Name	Type	Description	Default
`inputs`	`Optional[ndarray]`	Input data for the circuit. If `None`, default zero inputs are used.	`None`
`figure`	`str`	Rendering backend. One of: `"text"` - ASCII art (returned as a `str`). `"mpl"` - Matplotlib figure (returns `(fig, ax)`). `"tikz"` - LaTeX/TikZ `quantikz` code (returns a :class:`TikzFigure`). `"pulse"` - Pulse schedule (returns `(fig, axes)`). Only meaningful for pulse-mode models.	`'text'`
`**kwargs`	`Any`	Extra options forwarded to the drawing backend (e.g. `gate_values=True`).	`{}`

Returns:

Type	Description
`Union[str, Any]`	Depends on figure:
`Union[str, Any]`	`"text"` -> `str`
`Union[str, Any]`	`"mpl"` -> `(matplotlib.figure.Figure, matplotlib.axes.Axes)`
`Union[str, Any]`	`"tikz"` -> :class:`TikzFigure`

Raises:

Type	Description
`ValueError`	If figure is not one of the supported modes.

Source code in qml_essentials/model.py

def draw(
    self,
    inputs: Optional[jnp.ndarray] = None,
    figure: str = "text",
    **kwargs: Any,
) -> Union[str, Any]:
    """Visualize the quantum circuit.

    Records the circuit tape (without noise) and renders the gate
    sequence using the requested backend.

    Args:
        inputs (Optional[jnp.ndarray]): Input data for the circuit.
            If ``None``, default zero inputs are used.
        figure (str): Rendering backend.  One of:

            * ``"text"``  - ASCII art (returned as a ``str``).
            * ``"mpl"``   - Matplotlib figure (returns ``(fig, ax)``).
            * ``"tikz"``  - LaTeX/TikZ ``quantikz`` code (returns a
              :class:`TikzFigure`).
            * ``"pulse"`` - Pulse schedule (returns ``(fig, axes)``).
              Only meaningful for pulse-mode models.

        **kwargs: Extra options forwarded to the drawing backend
            (e.g. ``gate_values=True``).

    Returns:
        Depends on figure:

        * ``"text"``  -> ``str``
        * ``"mpl"``   -> ``(matplotlib.figure.Figure, matplotlib.axes.Axes)``
        * ``"tikz"``  -> :class:`TikzFigure`

    Raises:
        ValueError: If figure is not one of the supported modes.
    """
    inputs = self._inputs_validation(inputs)
    params = self.params[0] if self.params.ndim == 3 else self.params
    inp = inputs[0] if inputs.ndim == 2 else inputs

    if figure == "pulse":
        return self.draw_pulse(inputs=inputs, **kwargs)

    # Record without noise to get a clean circuit
    saved_noise = self._noise_params
    self._noise_params = None

    draw_script = js.Script(f=self._variational, n_qubits=self.n_qubits)
    result = draw_script.draw(
        figure=figure,
        args=(params, inp),
        kwargs={"noise_params": None},
        **kwargs,
    )

    self._noise_params = saved_noise
    return result

`draw_pulse(inputs=None, **kwargs)` #

Visualize the pulse schedule for the circuit.

Records the circuit in pulse mode and collects PulseEvents automatically via the pulse-event tape, then renders them.

State preparation, ansatz and encoding gates are all rendered as pulses. Encodings without a pulse parametrization (golomb and custom callables) are omitted from the schedule.

Parameters:

Name	Type	Description	Default
`inputs`	`Optional[ndarray]`	Input data. If `None`, default zero inputs are used.	`None`
`**kwargs`	`Any`	Forwarded to :func:`~jaqsi.drawing.draw_pulse_schedule` (e.g. `show_carrier=True`, `n_samples=300`).	`{}`

Returns:

Type	Description
`Any`	`(fig, axes)` — Matplotlib Figure and array of Axes.

Source code in qml_essentials/model.py

def draw_pulse(
    self,
    inputs: Optional[jnp.ndarray] = None,
    **kwargs: Any,
) -> Any:
    """Visualize the pulse schedule for the circuit.

    Records the circuit in pulse mode and collects PulseEvents
    automatically via the pulse-event tape, then renders them.

    State preparation, ansatz and encoding gates are all rendered as
    pulses. Encodings without a pulse parametrization (golomb and custom
    callables) are omitted from the schedule.

    Args:
        inputs: Input data.  If ``None``, default zero inputs are used.
        **kwargs: Forwarded to
            :func:`~jaqsi.drawing.draw_pulse_schedule`
            (e.g. ``show_carrier=True``, ``n_samples=300``).

    Returns:
        ``(fig, axes)`` — Matplotlib Figure and array of Axes.
    """
    if "gate_mode" in kwargs:
        warnings.warn(
            "draw_pulse no longer takes gate_mode, every gate group with a "
            "pulse representation is drawn.",
            DeprecationWarning,
            stacklevel=2,
        )
        kwargs.pop("gate_mode")

    inputs = self._inputs_validation(inputs)
    params = self.params[0] if self.params.ndim == 3 else self.params
    inp = inputs[0] if inputs.ndim == 2 else inputs

    # pass the model's own pulse parameters, so that _variational does not
    # fall back to them with a warning. Both are batch-first, so drawing
    # picks the first set, same as params above
    record_kwargs: Dict[str, Any] = {
        "gate_mode": "all_pulse" if self._enc_pulse_capable else "ansatz_pulse",
        "noise_params": None,
        "pulse_params": self.pulse_params[0],
    }
    if self._enc_pulse_capable:
        record_kwargs["enc_pulse_params"] = self.enc_pulse_params[0]

    draw_script = js.Script(f=self._variational, n_qubits=self.n_qubits)
    return draw_script.draw(
        figure="pulse",
        args=(params, inp),
        kwargs=record_kwargs,
        **kwargs,
    )

`exact_spectrum(method='tree')` #

Compute the exact per-feature Fourier spectrum via the FourierTree.

Unlike :attr:frequencies -- a naive per-feature estimate derived purely from the encoding, which can overestimate the spectrum (some coefficients are constrained to zero for all parameters) -- this builds the analytical Fourier tree (Nemkov et al.) and returns, for each input feature, the integer frequencies whose Fourier coefficient is not identically zero. The result is always a subset of :attr:frequencies.

The support is derived purely symbolically (no parameter sampling): see :meth:~qml_essentials.coefficients.FourierTree.get_exact_support. With method="tree" (default), frequencies whose contributions cancel identically across tree paths (e.g. two consecutive encodings combining into a single rotation) are excluded exactly; this enumerates the explicit tree, which can be infeasible for deep entangling circuits. With method="dp", a merged-state dynamic program derives the support without enumerating paths, which scales to deep circuits at the cost of not detecting identical cross-path cancellations.

Requires a Clifford + Pauli-rotation ansatz (see :class:~qml_essentials.pauli.PauliCircuit); other gate sets raise NotImplementedError during tree construction.

Parameters:

Name	Type	Description	Default
`method`	`str`	`"tree"` (fully exact) or `"dp"` (scalable).	`'tree'`

Returns:

Type	Description
`ndarray`	Tuple[np.ndarray, ...]: One sorted integer frequency array per input
`...`	feature (same layout as :attr:`frequencies`).

Source code in qml_essentials/model.py

def exact_spectrum(self, method: str = "tree") -> Tuple[np.ndarray, ...]:
    """Compute the exact per-feature Fourier spectrum via the FourierTree.

    Unlike :attr:`frequencies` -- a naive per-feature estimate derived purely
    from the encoding, which can *overestimate* the spectrum (some
    coefficients are constrained to zero for all parameters) -- this builds
    the analytical Fourier tree (Nemkov et al.) and returns, for each input
    feature, the integer frequencies whose Fourier coefficient is not
    identically zero.  The result is always a subset of :attr:`frequencies`.

    The support is derived purely symbolically (no parameter sampling): see
    :meth:`~qml_essentials.coefficients.FourierTree.get_exact_support`.
    With ``method="tree"`` (default), frequencies whose contributions cancel
    identically across tree paths (e.g. two consecutive encodings combining
    into a single rotation) are excluded exactly; this enumerates the
    explicit tree, which can be infeasible for deep entangling circuits.
    With ``method="dp"``, a merged-state dynamic program derives the support
    without enumerating paths, which scales to deep circuits at the cost of
    not detecting identical cross-path cancellations.

    Requires a Clifford + Pauli-rotation ansatz (see
    :class:`~qml_essentials.pauli.PauliCircuit`); other gate sets raise
    ``NotImplementedError`` during tree construction.

    Args:
        method (str): ``"tree"`` (fully exact) or ``"dp"`` (scalable).

    Returns:
        Tuple[np.ndarray, ...]: One sorted integer frequency array per input
        feature (same layout as :attr:`frequencies`).
    """
    from qml_essentials.coefficients import FourierTree  # avoid circular imp.

    tree = FourierTree(self)

    # Position of each model feature within the tree's frequency vectors.
    feature_pos = {feat: i for i, feat in enumerate(tree.features)}

    # Union of the symbolic supports over all observables (roots).
    support = set()
    for freqs in tree.get_exact_support(method=method):
        farr = np.asarray(freqs)
        for k in range(farr.shape[0]):
            key = (
                (int(farr[k]),)
                if farr.ndim == 1
                else tuple(int(v) for v in farr[k])
            )
            support.add(key)

    spectrum = []
    for feat in range(self.n_input_feat):
        if support and feat in feature_pos:
            pos = feature_pos[feat]
            vals = sorted({k[pos] for k in support})
        else:
            vals = [0]
        spectrum.append(np.array(vals, dtype=int))
    return tuple(spectrum)

`initialize_params(random_key=None, repeat=1, initialization=None, initialization_domain=None)` #

Initialize the variational parameters of the model.

Parameters:

Name	Type	Description	Default
`random_key`	`Optional[PRNGKey]`	JAX random key for initialization. If None, uses the model's internal random key.	`None`
`repeat`	`int`	Number of parameter sets to create (batch dimension). Defaults to 1.	`1`
`initialization`	`Optional[str]`	Strategy for parameter initialization. Options: "random", "zeros", "pi", "zero-controlled", "pi-controlled". If None, uses the strategy specified in the constructor.	`None`
`initialization_domain`	`Optional[List[float]]`	Domain [min, max] for random initialization. If None, uses the domain from constructor.	`None`

Returns:

Type	Description
`PRNGKey`	random.PRNGKey: Updated random key after initialization.

Raises:

Type	Description
`Exception`	If an invalid initialization method is specified.

Source code in qml_essentials/model.py

def initialize_params(
    self,
    random_key: Optional[random.PRNGKey] = None,
    repeat: int = 1,
    initialization: Optional[str] = None,
    initialization_domain: Optional[List[float]] = None,
) -> random.PRNGKey:
    """
    Initialize the variational parameters of the model.

    Args:
        random_key (Optional[random.PRNGKey]): JAX random key for initialization.
            If None, uses the model's internal random key.
        repeat (int): Number of parameter sets to create (batch dimension).
            Defaults to 1.
        initialization (Optional[str]): Strategy for parameter initialization.
            Options: "random", "zeros", "pi", "zero-controlled", "pi-controlled".
            If None, uses the strategy specified in the constructor.
        initialization_domain (Optional[List[float]]): Domain [min, max] for
            random initialization. If None, uses the domain from constructor.

    Returns:
        random.PRNGKey: Updated random key after initialization.

    Raises:
        Exception: If an invalid initialization method is specified.
    """
    # Initializing params
    params_shape = (repeat, *self._params_shape)

    # use existing strategy if not specified
    initialization = initialization or self._inialization_strategy
    initialization_domain = initialization_domain or self._initialization_domain

    random_key, sub_key = safe_random_split(
        random_key if random_key is not None else self.random_key
    )

    def set_control_params(params: jnp.ndarray, value: float) -> jnp.ndarray:
        indices = self.pqc.get_control_indices(self.n_qubits)
        if indices is None:
            warnings.warn(
                f"Specified {initialization} but circuit\
                does not contain controlled rotation gates.\
                Parameters are intialized randomly.",
                UserWarning,
            )
        else:
            np_params = np.array(params)
            np_params[:, :, indices[0] : indices[1] : indices[2]] = (
                np.ones_like(params[:, :, indices[0] : indices[1] : indices[2]])
                * value
            )
            params = jnp.array(np_params)
        return params

    if initialization == "random":
        self.params: jnp.ndarray = random.uniform(
            sub_key,
            params_shape,
            minval=initialization_domain[0],
            maxval=initialization_domain[1],
        )
    elif initialization == "zeros":
        self.params: jnp.ndarray = jnp.zeros(params_shape)
    elif initialization == "pi":
        self.params: jnp.ndarray = jnp.ones(params_shape) * jnp.pi
    elif initialization == "zero-controlled":
        self.params: jnp.ndarray = random.uniform(
            sub_key,
            params_shape,
            minval=initialization_domain[0],
            maxval=initialization_domain[1],
        )
        self.params = set_control_params(self.params, 0)
    elif initialization == "pi-controlled":
        self.params: jnp.ndarray = random.uniform(
            sub_key,
            params_shape,
            minval=initialization_domain[0],
            maxval=initialization_domain[1],
        )
        self.params = set_control_params(self.params, jnp.pi)
    else:
        raise Exception("Invalid initialization method")

    log.info(
        f"Initialized parameters with shape {self.params.shape}\
        using strategy {initialization}."
    )

    return random_key

`next_key()` #

Advance the internal random key and return a fresh sub key.

Intended for stochastic execution inside a JAX transform: a jitted call is traced once and replays the key that was current at trace time, so fresh randomness has to enter as an argument. Call this outside the transform and pass the result as random_key. Since the key is an argument rather than a constant, this does not trigger recompilation.

Returns:

Type	Description
`PRNGKey`	random.PRNGKey: Fresh sub key, split off the internal key.

Source code in qml_essentials/model.py

def next_key(self) -> random.PRNGKey:
    """
    Advance the internal random key and return a fresh sub key.

    Intended for stochastic execution inside a JAX transform: a jitted
    call is traced once and replays the key that was current at trace
    time, so fresh randomness has to enter as an argument. Call this
    outside the transform and pass the result as ``random_key``. Since the
    key is an argument rather than a constant, this does not trigger
    recompilation.

    Returns:
        random.PRNGKey: Fresh sub key, split off the internal key.
    """
    self.random_key, sub_key = safe_random_split(self.random_key)
    return sub_key

`transform_input(inputs, enc_params)` #

Transform input data by scaling with encoding parameters.

Implements the input transformation as described in arXiv:2309.03279v2, where inputs are linearly scaled by encoding parameters before being used in the quantum circuit.

Parameters:

Name	Type	Description	Default
`inputs`	`ndarray`	Input data point of shape (n_input_feat,) or (batch_size, n_input_feat).	required
`enc_params`	`ndarray`	Encoding weight scalar or vector used to scale the input.	required

Returns:

Type	Description
`ndarray`	jnp.ndarray: Transformed input, element-wise product of inputs and enc_params.

Source code in qml_essentials/model.py

def transform_input(
    self, inputs: jnp.ndarray, enc_params: jnp.ndarray
) -> jnp.ndarray:
    """
    Transform input data by scaling with encoding parameters.

    Implements the input transformation as described in arXiv:2309.03279v2,
    where inputs are linearly scaled by encoding parameters before being
    used in the quantum circuit.

    Args:
        inputs (jnp.ndarray): Input data point of shape (n_input_feat,) or
            (batch_size, n_input_feat).
        enc_params (jnp.ndarray): Encoding weight scalar or vector used to
            scale the input.

    Returns:
        jnp.ndarray: Transformed input, element-wise product of inputs
            and enc_params.
    """
    return inputs * enc_params

Entanglement#

from qml_essentials.entanglement import Entanglement

Source code in qml_essentials/entanglement.py

class Entanglement:
    @classmethod
    def meyer_wallach(
        cls,
        model: Model,
        n_samples: Optional[int | None],
        random_key: Optional[jax.random.PRNGKey] = None,
        scale: bool = False,
        **kwargs: Any,
    ) -> float:
        """
        Calculates the entangling capacity of a given quantum circuit
        using Meyer-Wallach measure.

        Args:
            model (Model): The quantum circuit model.
            n_samples (Optional[int]): Number of samples per qubit.
                If None or < 0, the current parameters of the model are used.
            random_key (Optional[jax.random.PRNGKey]): JAX random key for
                parameter initialization. If None, uses the model's internal
                random key.
            scale (bool): Whether to scale the number of samples.
            kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            float: Entangling capacity of the given circuit, guaranteed
                to be between 0.0 and 1.0.
        """
        if "noise_params" in kwargs:
            log.warning(
                "Meyer-Wallach measure not suitable for noisy circuits. "
                "Consider 'concentratable entanglement' instead."
            )

        if scale:
            n_samples = jnp.power(2, model.n_qubits) * n_samples

        if n_samples is not None and n_samples > 0:
            random_key = model.initialize_params(random_key, repeat=n_samples)

        # implicitly set input to none in case it's not needed
        kwargs.setdefault("inputs", None)
        # explicitly set execution type because everything else won't work
        rhos = model(execution_type="density", **kwargs).reshape(
            -1, 2**model.n_qubits, 2**model.n_qubits
        )

        ent = cls._compute_meyer_wallach_meas(rhos, model.n_qubits)

        log.debug(f"Variance of measure: {ent.var()}")

        return ent.mean()

    @classmethod
    def _compute_meyer_wallach_meas(
        cls, rhos: jnp.ndarray, n_qubits: int
    ) -> jnp.ndarray:
        """
        Computes the Meyer-Wallach entangling capability measure for a given
        set of density matrices.

        Args:
            rhos (jnp.ndarray): Density matrices of the sample quantum states.
                The shape is (B_s, 2^n, 2^n), where B_s is the number of samples
                (batch) and n the number of qubits
            n_qubits (int): The number of qubits

                    Returns:
            jnp.ndarray: Entangling capability for each sample, array with
                shape (B_s,)
        """
        qb = list(range(n_qubits))

        def _f(rhos):
            entropy = 0
            for j in range(n_qubits):
                # Formula 6 in https://doi.org/10.48550/arXiv.quant-ph/0305094
                # Trace out qubit j, keep all others
                keep = qb[:j] + qb[j + 1 :]
                density = js.partial_trace(rhos, n_qubits, keep)
                # only real values, because imaginary part will be separate
                # in all following calculations anyway
                # entropy should be 1/2 <= entropy <= 1
                entropy += jnp.trace((density @ density).real, axis1=-2, axis2=-1)

            # inverse averaged entropy and scale to [0, 1]
            return 2 * (1 - entropy / n_qubits)

        return jax.vmap(_f)(rhos)

    @classmethod
    def bell_measurements(
        cls,
        model: Model,
        n_samples: int,
        random_key: Optional[jax.random.PRNGKey] = None,
        scale: bool = False,
        **kwargs: Any,
    ) -> float:
        """
        Compute the Bell measurement for a given model.

        Constructs a ``2 * n_qubits`` circuit that prepares two copies of
        the model state (on disjoint qubit registers), applies CNOTs and
        Hadamards, and measures probabilities on the first register.

        Args:
            model (Model): The quantum circuit model.
            n_samples (int): The number of samples to compute the measure for.
            random_key (Optional[jax.random.PRNGKey]): JAX random key for
                parameter initialization. If None, uses the model's internal
                random key.
            scale (bool): Whether to scale the number of samples
                according to the number of qubits.
            **kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            float: The Bell measurement value.
        """
        if "noise_params" in kwargs:
            log.warning(
                "Bell Measurements not suitable for noisy circuits. "
                "Consider 'concentratable entanglement' instead."
            )

        if scale:
            n_samples = jnp.power(2, model.n_qubits) * n_samples

        n = model.n_qubits

        def _bell_circuit(params, inputs, pulse_params=None, random_key=None, **kw):
            """Bell measurement circuit on 2*n qubits."""
            from jaqsi.tape import copy_to_tape

            def vari():
                model._variational(
                    params,
                    inputs,
                    pulse_params=pulse_params,
                    random_key=random_key,
                    **kw,
                )

            # First copy on wires 0..n-1
            vari()
            # Second copy on wires n..2n-1
            copy_to_tape(vari, offset=n)

            # Bell measurement: CNOT + H
            for q in range(n):
                gateset.CX(wires=[q, q + n])
                gateset.H(wires=q)

        bell_script = js.Script(f=_bell_circuit, n_qubits=2 * n)

        if n_samples is not None and n_samples > 0:
            random_key = model.initialize_params(random_key, repeat=n_samples)
            params = model.params
        else:
            if len(model.params.shape) <= 2:
                params = model.params.reshape(1, *model.params.shape)
            else:
                log.info(f"Using sample size of model params: {model.params.shape[0]}")
                params = model.params

        n_samples = params.shape[0]
        inputs = model._inputs_validation(kwargs.get("inputs", None))

        # Execute: vmap over batch dimension of params (axis 0)
        if n_samples > 1:
            from jaqsi.utils import safe_random_split

            random_keys = safe_random_split(random_key, num=n_samples)
            result = bell_script.execute(
                type="probs",
                args=(params, inputs, model.pulse_params, random_keys),
                kwargs=kwargs,
                in_axes=(0, None, None, 0),
            )
        else:
            result = bell_script.execute(
                type="probs",
                args=(params, inputs, model.pulse_params, random_key),
                kwargs=kwargs,
            )

        # Marginalize: for each qubit q, keep wires [q, q+n] from the 2n-qubit probs
        # The last probability in each pair gives P(|11⟩) for that qubit pair
        per_qubit = []
        for q in range(n):
            marg = js.marginalize_probs(result, 2 * n, [q, q + n])
            per_qubit.append(marg)
        # per_qubit[q] has shape (n_samples, 4) or (4,)
        exp = jnp.stack(per_qubit, axis=-2)  # (..., n, 4)
        exp = 1 - 2 * exp[..., -1]  # (..., n)

        if not jnp.isclose(jnp.sum(exp.imag), 0, atol=1e-6):
            log.warning("Imaginary part of probabilities detected")
            exp = jnp.abs(exp)

        measure = 2 * (1 - exp.mean(axis=0))
        entangling_capability = min(max(float(measure.mean()), 0.0), 1.0)
        log.debug(f"Variance of measure: {measure.var()}")

        return entangling_capability

    @classmethod
    def relative_entropy(
        cls,
        model: Model,
        n_samples: int,
        n_sigmas: int,
        random_key: Optional[jax.random.PRNGKey] = None,
        scale: bool = False,
        **kwargs: Any,
    ) -> float:
        """
        Calculates the relative entropy of entanglement of a given quantum
        circuit. This measure is also applicable to mixed state, albeit it
        might me not fully accurate in this simplified case.

        As the relative entropy is generally defined as the smallest relative
        entropy from the state in question to the set of separable states.
        However, as computing the nearest separable state is NP-hard, we select
        n_sigmas of random separable states to compute the distance to, which
        is not necessarily the nearest. Thus, this measure of entanglement
        presents an upper limit of entanglement.

        As the relative entropy is not necessarily between zero and one, this
        function also normalises by the relative entroy to the GHZ state.

        Args:
            model (Model): The quantum circuit model.
            n_samples (int): Number of samples per qubit.
                If <= 0, the current parameters of the model are used.
            n_sigmas (int): Number of random separable pure states to compare against.
            random_key (Optional[jax.random.PRNGKey]): JAX random key for
                parameter initialization. If None, uses the model's internal
                random key.
            scale (bool): Whether to scale the number of samples.
            kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            float: Entangling capacity of the given circuit, guaranteed
                to be between 0.0 and 1.0.
        """
        dim = jnp.power(2, model.n_qubits)
        if scale:
            n_samples = dim * n_samples
            n_sigmas = dim * n_sigmas

        if random_key is None:
            random_key = model.random_key

        # Random separable states
        log_sigmas = sample_random_separable_states(
            model.n_qubits, n_samples=n_sigmas, random_key=random_key, take_log=True
        )

        random_key, _ = jax.random.split(random_key)

        if n_samples is not None and n_samples > 0:
            model.initialize_params(random_key, repeat=n_samples)
        else:
            if len(model.params.shape) <= 2:
                model.params = model.params.reshape(1, *model.params.shape)
            else:
                log.info(f"Using sample size of model params: {model.params.shape[0]}")

        rhos, log_rhos = cls._compute_log_density(model, **kwargs)

        rel_entropies = jnp.zeros((n_sigmas, model.params.shape[0]))

        for i, log_sigma in enumerate(log_sigmas):
            rel_entropies = rel_entropies.at[i].set(
                cls._compute_rel_entropies(rhos, log_rhos, log_sigma)
            )

        # Entropy of GHZ states should be maximal
        ghz_model = Model(model.n_qubits, 1, "GHZ", data_reupload=False)
        rho_ghz, log_rho_ghz = cls._compute_log_density(ghz_model, **kwargs)
        ghz_entropies = cls._compute_rel_entropies(rho_ghz, log_rho_ghz, log_sigmas)

        normalised_entropies = rel_entropies / ghz_entropies

        # Average all iterated states
        entangling_capability = normalised_entropies.T.min(axis=1)
        log.debug(f"Variance of measure: {entangling_capability.var()}")

        return entangling_capability.mean()

    @classmethod
    def _compute_log_density(
        cls, model: Model, **kwargs
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """
        Obtains the density matrix of a model and computes its logarithm.

        Args:
            model (Model): The model for which to compute the density matrix.

        Returns:
            Tuple[jnp.ndarray, jnp.ndarray]:
                - jnp.ndarray: density matrix.
                - jnp.ndarray: logarithm of the density matrix.
        """
        # implicitly set input to none in case it's not needed
        kwargs.setdefault("inputs", None)
        # explicitly set execution type because everything else won't work
        rho = model(execution_type="density", **kwargs)
        rho = rho.reshape(-1, 2**model.n_qubits, 2**model.n_qubits)
        log_rho = logm_v(rho) / jnp.log(2)
        return rho, log_rho

    @classmethod
    def _compute_rel_entropies(
        cls,
        rhos: jnp.ndarray,
        log_rhos: jnp.ndarray,
        log_sigmas: jnp.ndarray,
    ) -> jnp.ndarray:
        """
        Compute the relative entropy for a given model.

        Args:
            rhos (jnp.ndarray): Density matrix result of the circuit, has shape
                (R, 2^n, 2^n), with the batch size R and number of qubits n
            log_rhos (jnp.ndarray): Corresponding logarithm of the density
                matrix, has shape (R, 2^n, 2^n).
            log_sigmas (jnp.ndarray): Density matrix of next separable state,
                has shape (2^n, 2^n) if it's a single sigma or (S, 2^n, 2^n),
                with the batch size S (number of sigmas).

        Returns:
            jnp.ndarray: Relative Entropy for each sample
        """
        n_rhos = rhos.shape[0]
        if len(log_sigmas.shape) == 3:
            n_sigmas = log_sigmas.shape[0]
            rhos = jnp.tile(rhos, (n_sigmas, 1, 1))
            log_rhos = jnp.tile(log_rhos, (n_sigmas, 1, 1))
            einsum_subscript = "ij,jk->ik"
        else:
            n_sigmas = 1
            log_sigmas = log_sigmas[jnp.newaxis, ...].repeat(n_rhos, axis=0)

        einsum_subscript = "ij,jk->ik"

        def _f(rhos, log_rhos, log_sigmas):
            prod = jnp.einsum(einsum_subscript, rhos, log_rhos - log_sigmas)
            rel_entropies = jnp.abs(jnp.trace(prod, axis1=-2, axis2=-1))
            return rel_entropies

        rel_entropies = jax.vmap(_f, in_axes=(0, 0, 0))(rhos, log_rhos, log_sigmas)

        if n_sigmas > 1:
            rel_entropies = rel_entropies.reshape(n_sigmas, n_rhos)
        return rel_entropies

    @classmethod
    def entanglement_of_formation(
        cls,
        model: Model,
        n_samples: int,
        random_key: Optional[jax.random.PRNGKey] = None,
        scale: bool = False,
        always_decompose: bool = False,
        **kwargs: Any,
    ) -> float:
        """
        This function implements the entanglement of formation for mixed
        quantum systems.
        In that a mixed state gets decomposed into pure states with respective
        probabilities using the eigendecomposition of the density matrix.
        Then, the Meyer-Wallach measure is computed for each pure state,
        weighted by the eigenvalue.
        See e.g. https://doi.org/10.48550/arXiv.quant-ph/0504163

        Note that the decomposition is *not unique*! Therefore, this measure
        presents the entanglement for *some* decomposition into pure states,
        not necessarily the one that is anticipated when applying the Kraus
        channels.
        If a pure state is provided, this results in the same value as the
        Entanglement.meyer_wallach function if `always_decompose` flag is not set.

        Args:
            model (Model): The quantum circuit model.
            n_samples (int): Number of samples per qubit.
            random_key (Optional[jax.random.PRNGKey]): JAX random key for
                parameter initialization. If None, uses the model's internal
                random key.
            scale (bool): Whether to scale the number of samples.
            always_decompose (bool): Whether to explicitly compute the
                entantlement of formation for the eigendecomposition of a pure
                state.
            kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            float: Entangling capacity of the given circuit, guaranteed
                to be between 0.0 and 1.0.
        """

        if scale:
            n_samples = jnp.power(2, model.n_qubits) * n_samples

        if n_samples is not None and n_samples > 0:
            model.initialize_params(random_key, repeat=n_samples)
        else:
            if len(model.params.shape) <= 2:
                model.params = model.params.reshape(1, *model.params.shape)
            else:
                log.info(f"Using sample size of model params: {model.params.shape[0]}")

        # implicitly set input to none in case it's not needed
        kwargs.setdefault("inputs", None)
        rhos = model(execution_type="density", **kwargs)
        rhos = rhos.reshape(-1, 2**model.n_qubits, 2**model.n_qubits)
        ent = cls._compute_entanglement_of_formation(
            rhos, model.n_qubits, always_decompose
        )
        return ent.mean()

    @classmethod
    def _compute_entanglement_of_formation(
        cls,
        rhos: jnp.ndarray,
        n_qubits: int,
        always_decompose: bool,
    ) -> jnp.ndarray:
        """
        Computes the entanglement of formation for a given batch of density
        matrices.

        Args:
            rho (jnp.ndarray): The density matrices, has shape (B_s, 2^n, 2^n),
                where B_s is the batch size and n the number of qubits.
            n_qubits (int): Number of qubits
            always_decompose (bool): Whether to explicitly compute the
                entantlement of formation for the eigendecomposition of a pure
                state.

        Returns:
            jnp.ndarray: Entanglement for the provided density matrices.
        """
        eigenvalues, eigenvectors = jnp.linalg.eigh(rhos)
        if not always_decompose and jnp.isclose(eigenvalues, 1.0).any(axis=-1).all():
            return cls._compute_meyer_wallach_meas(rhos, n_qubits)

        rhos = np.einsum("sij,sik->sijk", eigenvectors, eigenvectors.conjugate())
        measures = cls._compute_meyer_wallach_meas(
            rhos.reshape(-1, 2**n_qubits, 2**n_qubits), n_qubits
        )
        ent = np.einsum("si,si->s", measures.reshape(-1, 2**n_qubits), eigenvalues)
        return ent

    @classmethod
    def concentratable_entanglement(
        cls,
        model: Model,
        n_samples: int,
        random_key: Optional[jax.random.PRNGKey] = None,
        scale: bool = False,
        **kwargs: Any,
    ) -> float:
        """
        Computes the concentratable entanglement of a given model.

        This method utilizes the Concentratable Entanglement measure from
        https://arxiv.org/abs/2104.06923.  The swap test is implemented
        directly in jaqsi using a ``3 * n_qubits`` circuit.

        Args:
            model (Model): The quantum circuit model.
            n_samples (int): The number of samples to compute the measure for.
            random_key (Optional[jax.random.PRNGKey]): JAX random key for
                parameter initialization. If None, uses the model's internal
                random key.
            scale (bool): Whether to scale the number of samples according to
                the number of qubits.
            **kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            float: Entangling capability of the given circuit, guaranteed
                to be between 0.0 and 1.0.
        """
        n = model.n_qubits
        N = 2**n

        if scale:
            n_samples = N * n_samples

        def _swap_test_circuit(
            params, inputs, pulse_params=None, random_key=None, **kw
        ):
            """Swap-test circuit on 3*n qubits."""
            from jaqsi.tape import copy_to_tape

            def vari():
                model._variational(
                    params,
                    inputs,
                    pulse_params=pulse_params,
                    random_key=random_key,
                    **kw,
                )

            # First copy on wires n..2n-1
            copy_to_tape(vari, offset=n)
            # Second copy on wires 2n..3n-1
            copy_to_tape(vari, offset=2 * n)

            # Swap test: H on ancilla register (wires 0..n-1)
            for i in range(n):
                gateset.H(wires=i)

            for i in range(n):
                gateset.CSWAP(wires=[i, i + n, i + 2 * n])

            for i in range(n):
                gateset.H(wires=i)

        swap_script = js.Script(f=_swap_test_circuit, n_qubits=3 * n)

        if n_samples is not None and n_samples > 0:
            random_key = model.initialize_params(random_key, repeat=n_samples)
        else:
            if len(model.params.shape) <= 2:
                model.params = model.params.reshape(1, *model.params.shape)
            else:
                log.info(f"Using sample size of model params: {model.params.shape[0]}")

        params = model.params
        inputs = model._inputs_validation(kwargs.get("inputs", None))
        n_batch = params.shape[0]

        marg_probs = jax.jit(js.marginalize_probs, static_argnums=(1, 2))

        if n_batch > 1:
            from jaqsi.utils import safe_random_split

            random_keys = safe_random_split(random_key, num=n_batch)
            probs = swap_script.execute(
                type="probs",
                args=(params, inputs, model.pulse_params, random_keys),
                in_axes=(0, None, None, 0),
                kwargs=kwargs,
            )
        else:
            probs = swap_script.execute(
                type="probs",
                args=(params, inputs, model.pulse_params, random_key),
                kwargs=kwargs,
            )

        # Marginalize to the ancilla register (wires 0..n-1)
        probs = marg_probs(probs, 3 * n, tuple(range(n)))

        ent = 1 - probs[..., 0]

        log.debug(f"Variance of measure: {ent.var()}")

        return float(ent.mean())

    @classmethod
    def concentratable_entanglement_estimation(
        cls,
        model: Model,
        n_samples: int,
        random_key: Optional[jax.random.PRNGKey] = None,
        scale: bool = False,
        **kwargs: Any,
    ) -> float:
        """
        Computes the concentratable entanglement of a given model.

        This method utilizes the Concentratable Entanglement measure from
        https://arxiv.org/abs/2104.06923.  The swap test is implemented
        directly in jaqsi using a ``3 * n_qubits`` circuit.

        Args:
            model (Model): The quantum circuit model.
            n_samples (int): The number of samples to compute the measure for.
            random_key (Optional[jax.random.PRNGKey]): JAX random key for
                parameter initialization. If None, uses the model's internal
                random key.
            scale (bool): Whether to scale the number of samples according to
                the number of qubits.
            **kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            float: Entangling capability of the given circuit, guaranteed
                to be between 0.0 and 1.0.
        """
        n = model.n_qubits
        N = 2**n

        if scale:
            n_samples = N * n_samples

        def _bell_basis_measurement(
            params, inputs, pulse_params=None, random_key=None, **kw
        ):
            """Bell-basis measurement circuit on 3*n qubits."""
            from jaqsi.tape import copy_to_tape

            def vari():
                model._variational(
                    params,
                    inputs,
                    pulse_params=pulse_params,
                    random_key=random_key,
                    **kw,
                )

            # First copy on wires 0..n-1
            copy_to_tape(vari, offset=0)
            # Second copy on wires n..2n-1
            copy_to_tape(vari, offset=n)

            for i in range(n):
                gateset.CX(wires=[i, i + n])
                gateset.H(wires=i)

        bell_basis_script = js.Script(f=_bell_basis_measurement, n_qubits=2 * n)

        if n_samples is not None and n_samples > 0:
            random_key = model.initialize_params(random_key, repeat=n_samples)
        else:
            if len(model.params.shape) <= 2:
                model.params = model.params.reshape(1, *model.params.shape)
            else:
                log.info(f"Using sample size of model params: {model.params.shape[0]}")

        params = model.params
        inputs = model._inputs_validation(kwargs.get("inputs", None))
        n_batch = params.shape[0]

        # SWAP operator in Bell-basis
        SWAP = jnp.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, -1]])
        # Construct observable for measuring CE
        CE_observable = gateset.Id([0, n]) + op.Operation([0, n], SWAP)
        for i in range(1, n):
            CE_observable = CE_observable @ (
                gateset.Id([i, i + n]) + op.Operation([i, i + n], SWAP)
            )
        CE_observable = (1 / N) * CE_observable

        expvals = []
        if n_batch > 1:
            from jaqsi.utils import safe_random_split

            random_keys = safe_random_split(random_key, num=n_batch)
            expvals = bell_basis_script.execute(
                type="expval",
                obs=[CE_observable],
                args=(params, inputs, model.pulse_params, random_keys),
                in_axes=(0, None, None, 0),
                kwargs=kwargs,
            )
        else:
            expvals = bell_basis_script.execute(
                type="expval",
                obs=[CE_observable],
                args=(params, inputs, model.pulse_params, random_key),
                kwargs=kwargs,
            )

        ent = 1 - expvals
        log.debug(f"Variance of measure: {ent.var()}")
        return float(ent.mean())

`bell_measurements(model, n_samples, random_key=None, scale=False, **kwargs)` `classmethod` #

Compute the Bell measurement for a given model.

Constructs a 2 * n_qubits circuit that prepares two copies of the model state (on disjoint qubit registers), applies CNOTs and Hadamards, and measures probabilities on the first register.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The quantum circuit model.	required
`n_samples`	`int`	The number of samples to compute the measure for.	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`scale`	`bool`	Whether to scale the number of samples according to the number of qubits.	`False`
`**kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Name	Type	Description
`float`	`float`	The Bell measurement value.

Source code in qml_essentials/entanglement.py

@classmethod
def bell_measurements(
    cls,
    model: Model,
    n_samples: int,
    random_key: Optional[jax.random.PRNGKey] = None,
    scale: bool = False,
    **kwargs: Any,
) -> float:
    """
    Compute the Bell measurement for a given model.

    Constructs a ``2 * n_qubits`` circuit that prepares two copies of
    the model state (on disjoint qubit registers), applies CNOTs and
    Hadamards, and measures probabilities on the first register.

    Args:
        model (Model): The quantum circuit model.
        n_samples (int): The number of samples to compute the measure for.
        random_key (Optional[jax.random.PRNGKey]): JAX random key for
            parameter initialization. If None, uses the model's internal
            random key.
        scale (bool): Whether to scale the number of samples
            according to the number of qubits.
        **kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        float: The Bell measurement value.
    """
    if "noise_params" in kwargs:
        log.warning(
            "Bell Measurements not suitable for noisy circuits. "
            "Consider 'concentratable entanglement' instead."
        )

    if scale:
        n_samples = jnp.power(2, model.n_qubits) * n_samples

    n = model.n_qubits

    def _bell_circuit(params, inputs, pulse_params=None, random_key=None, **kw):
        """Bell measurement circuit on 2*n qubits."""
        from jaqsi.tape import copy_to_tape

        def vari():
            model._variational(
                params,
                inputs,
                pulse_params=pulse_params,
                random_key=random_key,
                **kw,
            )

        # First copy on wires 0..n-1
        vari()
        # Second copy on wires n..2n-1
        copy_to_tape(vari, offset=n)

        # Bell measurement: CNOT + H
        for q in range(n):
            gateset.CX(wires=[q, q + n])
            gateset.H(wires=q)

    bell_script = js.Script(f=_bell_circuit, n_qubits=2 * n)

    if n_samples is not None and n_samples > 0:
        random_key = model.initialize_params(random_key, repeat=n_samples)
        params = model.params
    else:
        if len(model.params.shape) <= 2:
            params = model.params.reshape(1, *model.params.shape)
        else:
            log.info(f"Using sample size of model params: {model.params.shape[0]}")
            params = model.params

    n_samples = params.shape[0]
    inputs = model._inputs_validation(kwargs.get("inputs", None))

    # Execute: vmap over batch dimension of params (axis 0)
    if n_samples > 1:
        from jaqsi.utils import safe_random_split

        random_keys = safe_random_split(random_key, num=n_samples)
        result = bell_script.execute(
            type="probs",
            args=(params, inputs, model.pulse_params, random_keys),
            kwargs=kwargs,
            in_axes=(0, None, None, 0),
        )
    else:
        result = bell_script.execute(
            type="probs",
            args=(params, inputs, model.pulse_params, random_key),
            kwargs=kwargs,
        )

    # Marginalize: for each qubit q, keep wires [q, q+n] from the 2n-qubit probs
    # The last probability in each pair gives P(|11⟩) for that qubit pair
    per_qubit = []
    for q in range(n):
        marg = js.marginalize_probs(result, 2 * n, [q, q + n])
        per_qubit.append(marg)
    # per_qubit[q] has shape (n_samples, 4) or (4,)
    exp = jnp.stack(per_qubit, axis=-2)  # (..., n, 4)
    exp = 1 - 2 * exp[..., -1]  # (..., n)

    if not jnp.isclose(jnp.sum(exp.imag), 0, atol=1e-6):
        log.warning("Imaginary part of probabilities detected")
        exp = jnp.abs(exp)

    measure = 2 * (1 - exp.mean(axis=0))
    entangling_capability = min(max(float(measure.mean()), 0.0), 1.0)
    log.debug(f"Variance of measure: {measure.var()}")

    return entangling_capability

`concentratable_entanglement(model, n_samples, random_key=None, scale=False, **kwargs)` `classmethod` #

Computes the concentratable entanglement of a given model.

This method utilizes the Concentratable Entanglement measure from https://arxiv.org/abs/2104.06923. The swap test is implemented directly in jaqsi using a 3 * n_qubits circuit.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The quantum circuit model.	required
`n_samples`	`int`	The number of samples to compute the measure for.	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`scale`	`bool`	Whether to scale the number of samples according to the number of qubits.	`False`
`**kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Name	Type	Description
`float`	`float`	Entangling capability of the given circuit, guaranteed to be between 0.0 and 1.0.

Source code in qml_essentials/entanglement.py

@classmethod
def concentratable_entanglement(
    cls,
    model: Model,
    n_samples: int,
    random_key: Optional[jax.random.PRNGKey] = None,
    scale: bool = False,
    **kwargs: Any,
) -> float:
    """
    Computes the concentratable entanglement of a given model.

    This method utilizes the Concentratable Entanglement measure from
    https://arxiv.org/abs/2104.06923.  The swap test is implemented
    directly in jaqsi using a ``3 * n_qubits`` circuit.

    Args:
        model (Model): The quantum circuit model.
        n_samples (int): The number of samples to compute the measure for.
        random_key (Optional[jax.random.PRNGKey]): JAX random key for
            parameter initialization. If None, uses the model's internal
            random key.
        scale (bool): Whether to scale the number of samples according to
            the number of qubits.
        **kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        float: Entangling capability of the given circuit, guaranteed
            to be between 0.0 and 1.0.
    """
    n = model.n_qubits
    N = 2**n

    if scale:
        n_samples = N * n_samples

    def _swap_test_circuit(
        params, inputs, pulse_params=None, random_key=None, **kw
    ):
        """Swap-test circuit on 3*n qubits."""
        from jaqsi.tape import copy_to_tape

        def vari():
            model._variational(
                params,
                inputs,
                pulse_params=pulse_params,
                random_key=random_key,
                **kw,
            )

        # First copy on wires n..2n-1
        copy_to_tape(vari, offset=n)
        # Second copy on wires 2n..3n-1
        copy_to_tape(vari, offset=2 * n)

        # Swap test: H on ancilla register (wires 0..n-1)
        for i in range(n):
            gateset.H(wires=i)

        for i in range(n):
            gateset.CSWAP(wires=[i, i + n, i + 2 * n])

        for i in range(n):
            gateset.H(wires=i)

    swap_script = js.Script(f=_swap_test_circuit, n_qubits=3 * n)

    if n_samples is not None and n_samples > 0:
        random_key = model.initialize_params(random_key, repeat=n_samples)
    else:
        if len(model.params.shape) <= 2:
            model.params = model.params.reshape(1, *model.params.shape)
        else:
            log.info(f"Using sample size of model params: {model.params.shape[0]}")

    params = model.params
    inputs = model._inputs_validation(kwargs.get("inputs", None))
    n_batch = params.shape[0]

    marg_probs = jax.jit(js.marginalize_probs, static_argnums=(1, 2))

    if n_batch > 1:
        from jaqsi.utils import safe_random_split

        random_keys = safe_random_split(random_key, num=n_batch)
        probs = swap_script.execute(
            type="probs",
            args=(params, inputs, model.pulse_params, random_keys),
            in_axes=(0, None, None, 0),
            kwargs=kwargs,
        )
    else:
        probs = swap_script.execute(
            type="probs",
            args=(params, inputs, model.pulse_params, random_key),
            kwargs=kwargs,
        )

    # Marginalize to the ancilla register (wires 0..n-1)
    probs = marg_probs(probs, 3 * n, tuple(range(n)))

    ent = 1 - probs[..., 0]

    log.debug(f"Variance of measure: {ent.var()}")

    return float(ent.mean())

`concentratable_entanglement_estimation(model, n_samples, random_key=None, scale=False, **kwargs)` `classmethod` #

Computes the concentratable entanglement of a given model.

This method utilizes the Concentratable Entanglement measure from https://arxiv.org/abs/2104.06923. The swap test is implemented directly in jaqsi using a 3 * n_qubits circuit.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The quantum circuit model.	required
`n_samples`	`int`	The number of samples to compute the measure for.	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`scale`	`bool`	Whether to scale the number of samples according to the number of qubits.	`False`
`**kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Name	Type	Description
`float`	`float`	Entangling capability of the given circuit, guaranteed to be between 0.0 and 1.0.

Source code in qml_essentials/entanglement.py

@classmethod
def concentratable_entanglement_estimation(
    cls,
    model: Model,
    n_samples: int,
    random_key: Optional[jax.random.PRNGKey] = None,
    scale: bool = False,
    **kwargs: Any,
) -> float:
    """
    Computes the concentratable entanglement of a given model.

    This method utilizes the Concentratable Entanglement measure from
    https://arxiv.org/abs/2104.06923.  The swap test is implemented
    directly in jaqsi using a ``3 * n_qubits`` circuit.

    Args:
        model (Model): The quantum circuit model.
        n_samples (int): The number of samples to compute the measure for.
        random_key (Optional[jax.random.PRNGKey]): JAX random key for
            parameter initialization. If None, uses the model's internal
            random key.
        scale (bool): Whether to scale the number of samples according to
            the number of qubits.
        **kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        float: Entangling capability of the given circuit, guaranteed
            to be between 0.0 and 1.0.
    """
    n = model.n_qubits
    N = 2**n

    if scale:
        n_samples = N * n_samples

    def _bell_basis_measurement(
        params, inputs, pulse_params=None, random_key=None, **kw
    ):
        """Bell-basis measurement circuit on 3*n qubits."""
        from jaqsi.tape import copy_to_tape

        def vari():
            model._variational(
                params,
                inputs,
                pulse_params=pulse_params,
                random_key=random_key,
                **kw,
            )

        # First copy on wires 0..n-1
        copy_to_tape(vari, offset=0)
        # Second copy on wires n..2n-1
        copy_to_tape(vari, offset=n)

        for i in range(n):
            gateset.CX(wires=[i, i + n])
            gateset.H(wires=i)

    bell_basis_script = js.Script(f=_bell_basis_measurement, n_qubits=2 * n)

    if n_samples is not None and n_samples > 0:
        random_key = model.initialize_params(random_key, repeat=n_samples)
    else:
        if len(model.params.shape) <= 2:
            model.params = model.params.reshape(1, *model.params.shape)
        else:
            log.info(f"Using sample size of model params: {model.params.shape[0]}")

    params = model.params
    inputs = model._inputs_validation(kwargs.get("inputs", None))
    n_batch = params.shape[0]

    # SWAP operator in Bell-basis
    SWAP = jnp.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, -1]])
    # Construct observable for measuring CE
    CE_observable = gateset.Id([0, n]) + op.Operation([0, n], SWAP)
    for i in range(1, n):
        CE_observable = CE_observable @ (
            gateset.Id([i, i + n]) + op.Operation([i, i + n], SWAP)
        )
    CE_observable = (1 / N) * CE_observable

    expvals = []
    if n_batch > 1:
        from jaqsi.utils import safe_random_split

        random_keys = safe_random_split(random_key, num=n_batch)
        expvals = bell_basis_script.execute(
            type="expval",
            obs=[CE_observable],
            args=(params, inputs, model.pulse_params, random_keys),
            in_axes=(0, None, None, 0),
            kwargs=kwargs,
        )
    else:
        expvals = bell_basis_script.execute(
            type="expval",
            obs=[CE_observable],
            args=(params, inputs, model.pulse_params, random_key),
            kwargs=kwargs,
        )

    ent = 1 - expvals
    log.debug(f"Variance of measure: {ent.var()}")
    return float(ent.mean())

`entanglement_of_formation(model, n_samples, random_key=None, scale=False, always_decompose=False, **kwargs)` `classmethod` #

This function implements the entanglement of formation for mixed quantum systems. In that a mixed state gets decomposed into pure states with respective probabilities using the eigendecomposition of the density matrix. Then, the Meyer-Wallach measure is computed for each pure state, weighted by the eigenvalue. See e.g. https://doi.org/10.48550/arXiv.quant-ph/0504163

Note that the decomposition is not unique! Therefore, this measure presents the entanglement for some decomposition into pure states, not necessarily the one that is anticipated when applying the Kraus channels. If a pure state is provided, this results in the same value as the Entanglement.meyer_wallach function if always_decompose flag is not set.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The quantum circuit model.	required
`n_samples`	`int`	Number of samples per qubit.	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`scale`	`bool`	Whether to scale the number of samples.	`False`
`always_decompose`	`bool`	Whether to explicitly compute the entantlement of formation for the eigendecomposition of a pure state.	`False`
`kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Name	Type	Description
`float`	`float`	Entangling capacity of the given circuit, guaranteed to be between 0.0 and 1.0.

Source code in qml_essentials/entanglement.py

@classmethod
def entanglement_of_formation(
    cls,
    model: Model,
    n_samples: int,
    random_key: Optional[jax.random.PRNGKey] = None,
    scale: bool = False,
    always_decompose: bool = False,
    **kwargs: Any,
) -> float:
    """
    This function implements the entanglement of formation for mixed
    quantum systems.
    In that a mixed state gets decomposed into pure states with respective
    probabilities using the eigendecomposition of the density matrix.
    Then, the Meyer-Wallach measure is computed for each pure state,
    weighted by the eigenvalue.
    See e.g. https://doi.org/10.48550/arXiv.quant-ph/0504163

    Note that the decomposition is *not unique*! Therefore, this measure
    presents the entanglement for *some* decomposition into pure states,
    not necessarily the one that is anticipated when applying the Kraus
    channels.
    If a pure state is provided, this results in the same value as the
    Entanglement.meyer_wallach function if `always_decompose` flag is not set.

    Args:
        model (Model): The quantum circuit model.
        n_samples (int): Number of samples per qubit.
        random_key (Optional[jax.random.PRNGKey]): JAX random key for
            parameter initialization. If None, uses the model's internal
            random key.
        scale (bool): Whether to scale the number of samples.
        always_decompose (bool): Whether to explicitly compute the
            entantlement of formation for the eigendecomposition of a pure
            state.
        kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        float: Entangling capacity of the given circuit, guaranteed
            to be between 0.0 and 1.0.
    """

    if scale:
        n_samples = jnp.power(2, model.n_qubits) * n_samples

    if n_samples is not None and n_samples > 0:
        model.initialize_params(random_key, repeat=n_samples)
    else:
        if len(model.params.shape) <= 2:
            model.params = model.params.reshape(1, *model.params.shape)
        else:
            log.info(f"Using sample size of model params: {model.params.shape[0]}")

    # implicitly set input to none in case it's not needed
    kwargs.setdefault("inputs", None)
    rhos = model(execution_type="density", **kwargs)
    rhos = rhos.reshape(-1, 2**model.n_qubits, 2**model.n_qubits)
    ent = cls._compute_entanglement_of_formation(
        rhos, model.n_qubits, always_decompose
    )
    return ent.mean()

`meyer_wallach(model, n_samples, random_key=None, scale=False, **kwargs)` `classmethod` #

Calculates the entangling capacity of a given quantum circuit using Meyer-Wallach measure.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The quantum circuit model.	required
`n_samples`	`Optional[int]`	Number of samples per qubit. If None or < 0, the current parameters of the model are used.	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`scale`	`bool`	Whether to scale the number of samples.	`False`
`kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Name	Type	Description
`float`	`float`	Entangling capacity of the given circuit, guaranteed to be between 0.0 and 1.0.

Source code in qml_essentials/entanglement.py

@classmethod
def meyer_wallach(
    cls,
    model: Model,
    n_samples: Optional[int | None],
    random_key: Optional[jax.random.PRNGKey] = None,
    scale: bool = False,
    **kwargs: Any,
) -> float:
    """
    Calculates the entangling capacity of a given quantum circuit
    using Meyer-Wallach measure.

    Args:
        model (Model): The quantum circuit model.
        n_samples (Optional[int]): Number of samples per qubit.
            If None or < 0, the current parameters of the model are used.
        random_key (Optional[jax.random.PRNGKey]): JAX random key for
            parameter initialization. If None, uses the model's internal
            random key.
        scale (bool): Whether to scale the number of samples.
        kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        float: Entangling capacity of the given circuit, guaranteed
            to be between 0.0 and 1.0.
    """
    if "noise_params" in kwargs:
        log.warning(
            "Meyer-Wallach measure not suitable for noisy circuits. "
            "Consider 'concentratable entanglement' instead."
        )

    if scale:
        n_samples = jnp.power(2, model.n_qubits) * n_samples

    if n_samples is not None and n_samples > 0:
        random_key = model.initialize_params(random_key, repeat=n_samples)

    # implicitly set input to none in case it's not needed
    kwargs.setdefault("inputs", None)
    # explicitly set execution type because everything else won't work
    rhos = model(execution_type="density", **kwargs).reshape(
        -1, 2**model.n_qubits, 2**model.n_qubits
    )

    ent = cls._compute_meyer_wallach_meas(rhos, model.n_qubits)

    log.debug(f"Variance of measure: {ent.var()}")

    return ent.mean()

`relative_entropy(model, n_samples, n_sigmas, random_key=None, scale=False, **kwargs)` `classmethod` #

Calculates the relative entropy of entanglement of a given quantum circuit. This measure is also applicable to mixed state, albeit it might me not fully accurate in this simplified case.

As the relative entropy is generally defined as the smallest relative entropy from the state in question to the set of separable states. However, as computing the nearest separable state is NP-hard, we select n_sigmas of random separable states to compute the distance to, which is not necessarily the nearest. Thus, this measure of entanglement presents an upper limit of entanglement.

As the relative entropy is not necessarily between zero and one, this function also normalises by the relative entroy to the GHZ state.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The quantum circuit model.	required
`n_samples`	`int`	Number of samples per qubit. If <= 0, the current parameters of the model are used.	required
`n_sigmas`	`int`	Number of random separable pure states to compare against.	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`scale`	`bool`	Whether to scale the number of samples.	`False`
`kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Name	Type	Description
`float`	`float`	Entangling capacity of the given circuit, guaranteed to be between 0.0 and 1.0.

Source code in qml_essentials/entanglement.py

@classmethod
def relative_entropy(
    cls,
    model: Model,
    n_samples: int,
    n_sigmas: int,
    random_key: Optional[jax.random.PRNGKey] = None,
    scale: bool = False,
    **kwargs: Any,
) -> float:
    """
    Calculates the relative entropy of entanglement of a given quantum
    circuit. This measure is also applicable to mixed state, albeit it
    might me not fully accurate in this simplified case.

    As the relative entropy is generally defined as the smallest relative
    entropy from the state in question to the set of separable states.
    However, as computing the nearest separable state is NP-hard, we select
    n_sigmas of random separable states to compute the distance to, which
    is not necessarily the nearest. Thus, this measure of entanglement
    presents an upper limit of entanglement.

    As the relative entropy is not necessarily between zero and one, this
    function also normalises by the relative entroy to the GHZ state.

    Args:
        model (Model): The quantum circuit model.
        n_samples (int): Number of samples per qubit.
            If <= 0, the current parameters of the model are used.
        n_sigmas (int): Number of random separable pure states to compare against.
        random_key (Optional[jax.random.PRNGKey]): JAX random key for
            parameter initialization. If None, uses the model's internal
            random key.
        scale (bool): Whether to scale the number of samples.
        kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        float: Entangling capacity of the given circuit, guaranteed
            to be between 0.0 and 1.0.
    """
    dim = jnp.power(2, model.n_qubits)
    if scale:
        n_samples = dim * n_samples
        n_sigmas = dim * n_sigmas

    if random_key is None:
        random_key = model.random_key

    # Random separable states
    log_sigmas = sample_random_separable_states(
        model.n_qubits, n_samples=n_sigmas, random_key=random_key, take_log=True
    )

    random_key, _ = jax.random.split(random_key)

    if n_samples is not None and n_samples > 0:
        model.initialize_params(random_key, repeat=n_samples)
    else:
        if len(model.params.shape) <= 2:
            model.params = model.params.reshape(1, *model.params.shape)
        else:
            log.info(f"Using sample size of model params: {model.params.shape[0]}")

    rhos, log_rhos = cls._compute_log_density(model, **kwargs)

    rel_entropies = jnp.zeros((n_sigmas, model.params.shape[0]))

    for i, log_sigma in enumerate(log_sigmas):
        rel_entropies = rel_entropies.at[i].set(
            cls._compute_rel_entropies(rhos, log_rhos, log_sigma)
        )

    # Entropy of GHZ states should be maximal
    ghz_model = Model(model.n_qubits, 1, "GHZ", data_reupload=False)
    rho_ghz, log_rho_ghz = cls._compute_log_density(ghz_model, **kwargs)
    ghz_entropies = cls._compute_rel_entropies(rho_ghz, log_rho_ghz, log_sigmas)

    normalised_entropies = rel_entropies / ghz_entropies

    # Average all iterated states
    entangling_capability = normalised_entropies.T.min(axis=1)
    log.debug(f"Variance of measure: {entangling_capability.var()}")

    return entangling_capability.mean()

Expressibility#

from qml_essentials.expressibility import Expressibility

Source code in qml_essentials/expressibility.py

class Expressibility:
    @classmethod
    def _sample_state_fidelities(
        cls,
        model: Model,
        n_samples: int,
        random_key: Optional[jax.random.PRNGKey] = None,
        kwargs: Any = None,
    ) -> jnp.ndarray:
        """
        Compute the fidelities for each parameter set.

        Args:
            model (Callable): Function that models the quantum circuit.
            n_samples (int): Number of parameter sets to generate.
            random_key (Optional[jax.random.PRNGKey]): JAX random key for
                parameter initialization. If None, uses the model's internal
                random key.
            kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            jnp.ndarray: Array of shape (n_samples,) containing the fidelities.
        """
        # Generate random parameter sets
        # We need two sets of parameters, as we are computing fidelities for a
        # pair of random state vectors
        model.initialize_params(random_key, repeat=n_samples * 2)

        # Evaluate the model for all parameters
        # Execution type is explicitly set to density
        sv: jnp.ndarray = model(
            params=model.params,
            execution_type="density",
            **kwargs,
        )

        # $\sqrt{\rho}$
        sqrt_sv1: jnp.ndarray = jnp.array([sqrtm(m) for m in sv[:n_samples]])

        # $\sqrt{\rho} \sigma \sqrt{\rho}$
        inner_fidelity = sqrt_sv1 @ sv[n_samples:] @ sqrt_sv1

        # Compute the fidelity using the partial trace of the statevector
        fidelity: jnp.ndarray = (
            jnp.trace(
                jnp.array([sqrtm(m) for m in inner_fidelity]),
                axis1=1,
                axis2=2,
            )
            ** 2
        )

        fidelity = jnp.abs(fidelity)

        return fidelity

    @classmethod
    def state_fidelities(
        cls,
        n_samples: int,
        n_bins: int,
        model: Model,
        random_key: Optional[jax.random.PRNGKey] = None,
        scale: bool = False,
        **kwargs: Any,
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """
        Sample the state fidelities and histogram them into a 2D array.

        Args:
            n_samples (int): Number of parameter sets to generate.
            n_bins (int): Number of histogram bins.
            model (Callable): Function that models the quantum circuit.
            random_key (Optional[jax.random.PRNGKey]): JAX random key for
                parameter initialization. If None, uses the model's internal
                random key.
            scale (bool): Whether to scale the number of samples and bins.
            kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            Tuple[jnp.ndarray, jnp.ndarray]: Tuple containing the bin edges,
            and histogram values.
        """
        if scale:
            n_samples = jnp.power(2, model.n_qubits) * n_samples
            n_bins = model.n_qubits * n_bins

        fidelities = cls._sample_state_fidelities(
            n_samples=n_samples,
            random_key=random_key,
            model=model,
            kwargs=kwargs,
        )

        y: jnp.ndarray = jnp.linspace(0, 1, n_bins + 1)

        z, _ = jnp.histogram(fidelities, bins=y)

        z = z / n_samples

        return y, z

    @classmethod
    def _haar_probability(cls, fidelity: float, n_qubits: int) -> float:
        """
        Calculates theoretical probability density function for random Haar states
        as proposed by Sim et al. (https://arxiv.org/abs/1905.10876).

        Args:
            fidelity (float): fidelity of two parameter assignments in [0, 1]
            n_qubits (int): number of qubits in the quantum system

        Returns:
            float: probability for a given fidelity
        """
        N = 2**n_qubits

        prob = (N - 1) * (1 - fidelity) ** (N - 2)
        return prob

    @classmethod
    def _sample_haar_integral(cls, n_qubits: int, n_bins: int) -> jnp.ndarray:
        """
        Calculates theoretical probability density function for random Haar states
        as proposed by Sim et al. (https://arxiv.org/abs/1905.10876) and bins it
        into a 2D-histogram.

        Args:
            n_qubits (int): number of qubits in the quantum system
            n_bins (int): number of histogram bins

        Returns:
            jnp.ndarray: probability distribution for all fidelities
        """
        dist = np.zeros(n_bins)
        for idx in range(n_bins):
            v = idx / n_bins
            u = (idx + 1) / n_bins
            dist[idx], _ = integrate.quad(cls._haar_probability, v, u, args=(n_qubits,))

        return dist

    @classmethod
    def haar_integral(
        cls,
        n_qubits: int,
        n_bins: int,
        cache: bool = True,
        scale: bool = False,
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """
        Calculates theoretical probability density function for random Haar states
        as proposed by Sim et al. (https://arxiv.org/abs/1905.10876) and bins it
        into a 3D-histogram.

        Args:
            n_qubits (int): number of qubits in the quantum system
            n_bins (int): number of histogram bins
            cache (bool): whether to cache the haar integral
            scale (bool): whether to scale the number of bins

        Returns:
            Tuple[jnp.ndarray, jnp.ndarray]:
                - x component (bins): the input domain
                - y component (probabilities): the haar probability density
                  funtion for random Haar states
        """
        if scale:
            n_bins = n_qubits * n_bins

        x = jnp.linspace(0, 1, n_bins)

        if cache:
            name = f"haar_{n_qubits}q_{n_bins}s_{'scaled' if scale else ''}.npy"

            cache_folder = ".cache"
            if not os.path.exists(cache_folder):
                os.mkdir(cache_folder)

            file_path = os.path.join(cache_folder, name)

            if os.path.isfile(file_path):
                y = jnp.load(file_path)
                return x, y

        y = cls._sample_haar_integral(n_qubits, n_bins)

        if cache:
            jnp.save(file_path, y)

        return x, y

    @classmethod
    def kullback_leibler_divergence(
        cls,
        vqc_prob_dist: jnp.ndarray,
        haar_dist: jnp.ndarray,
    ) -> jnp.ndarray:
        """
        Calculates the KL divergence between two probability distributions (Haar
        probability distribution and the fidelity distribution sampled from a VQC).

        Args:
            vqc_prob_dist (jnp.ndarray): VQC fidelity probability distribution.
                Should have shape (n_inputs_samples, n_bins)
            haar_dist (jnp.ndarray): Haar probability distribution with shape.
                Should have shape (n_bins, )

        Returns:
            jnp.ndarray: Array of KL-Divergence values for all values in axis 1
        """
        if len(vqc_prob_dist.shape) > 1:
            assert all([haar_dist.shape == p.shape for p in vqc_prob_dist]), (
                "All probabilities for inputs should have the same shape as Haar. "
                f"Got {haar_dist.shape} for Haar and {vqc_prob_dist.shape} for VQC"
            )
        else:
            vqc_prob_dist = vqc_prob_dist.reshape((1, -1))

        kl_divergence = np.zeros(vqc_prob_dist.shape[0])
        for idx, p in enumerate(vqc_prob_dist):
            kl_divergence[idx] = jnp.sum(rel_entr(p, haar_dist))

        return kl_divergence

    @classmethod
    def kl_divergence_to_haar(
        cls,
        model: Model,
        n_samples: int,
        n_bins: int,
        random_key: Optional[jax.random.PRNGKey] = None,
        scale: bool = False,
        **kwargs: Any,
    ) -> float:
        """
        Shortcut method to compute the KL-Divergence bewteen a model and the
        Haar distribution. The basic steps are:
            - Sample the state fidelities for randomly initialised parameters.
            - Calculates the KL divergence between the sampled probability and
              the Haar probability distribution.

        Args:
            model (Model): Function that models the quantum circuit.
            n_samples (int): Number of parameter sets to generate.
            n_bins (int): Number of histogram bins.
            random_key (Optional[jax.random.PRNGKey]): JAX random key for
                parameter initialization. If None, uses the model's internal
                random key.
            scale (bool): Whether to scale the number of samples and bins.
            kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: Tuple containing the
                input samples, bin edges, and histogram values.
        """
        _, fidelities = Expressibility.state_fidelities(
            model=model,
            random_key=random_key,
            n_samples=n_samples,
            n_bins=n_bins,
            scale=scale,
            **kwargs,
        )
        _, haar_probs = Expressibility.haar_integral(
            model.n_qubits, n_bins=n_bins, scale=scale
        )
        return Expressibility.kullback_leibler_divergence(fidelities, haar_probs)

`haar_integral(n_qubits, n_bins, cache=True, scale=False)` `classmethod` #

Calculates theoretical probability density function for random Haar states as proposed by Sim et al. (https://arxiv.org/abs/1905.10876) and bins it into a 3D-histogram.

Parameters:

Name	Type	Description	Default
`n_qubits`	`int`	number of qubits in the quantum system	required
`n_bins`	`int`	number of histogram bins	required
`cache`	`bool`	whether to cache the haar integral	`True`
`scale`	`bool`	whether to scale the number of bins	`False`

Returns:

Type	Description
`Tuple[ndarray, ndarray]`	Tuple[jnp.ndarray, jnp.ndarray]: - x component (bins): the input domain - y component (probabilities): the haar probability density funtion for random Haar states

Source code in qml_essentials/expressibility.py

@classmethod
def haar_integral(
    cls,
    n_qubits: int,
    n_bins: int,
    cache: bool = True,
    scale: bool = False,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """
    Calculates theoretical probability density function for random Haar states
    as proposed by Sim et al. (https://arxiv.org/abs/1905.10876) and bins it
    into a 3D-histogram.

    Args:
        n_qubits (int): number of qubits in the quantum system
        n_bins (int): number of histogram bins
        cache (bool): whether to cache the haar integral
        scale (bool): whether to scale the number of bins

    Returns:
        Tuple[jnp.ndarray, jnp.ndarray]:
            - x component (bins): the input domain
            - y component (probabilities): the haar probability density
              funtion for random Haar states
    """
    if scale:
        n_bins = n_qubits * n_bins

    x = jnp.linspace(0, 1, n_bins)

    if cache:
        name = f"haar_{n_qubits}q_{n_bins}s_{'scaled' if scale else ''}.npy"

        cache_folder = ".cache"
        if not os.path.exists(cache_folder):
            os.mkdir(cache_folder)

        file_path = os.path.join(cache_folder, name)

        if os.path.isfile(file_path):
            y = jnp.load(file_path)
            return x, y

    y = cls._sample_haar_integral(n_qubits, n_bins)

    if cache:
        jnp.save(file_path, y)

    return x, y

`kl_divergence_to_haar(model, n_samples, n_bins, random_key=None, scale=False, **kwargs)` `classmethod` #

Shortcut method to compute the KL-Divergence bewteen a model and the Haar distribution. The basic steps are: - Sample the state fidelities for randomly initialised parameters. - Calculates the KL divergence between the sampled probability and the Haar probability distribution.

Parameters:

Name	Type	Description	Default
`model`	`Model`	Function that models the quantum circuit.	required
`n_samples`	`int`	Number of parameter sets to generate.	required
`n_bins`	`int`	Number of histogram bins.	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`scale`	`bool`	Whether to scale the number of samples and bins.	`False`
`kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Type	Description
`float`	Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: Tuple containing the input samples, bin edges, and histogram values.

Source code in qml_essentials/expressibility.py

@classmethod
def kl_divergence_to_haar(
    cls,
    model: Model,
    n_samples: int,
    n_bins: int,
    random_key: Optional[jax.random.PRNGKey] = None,
    scale: bool = False,
    **kwargs: Any,
) -> float:
    """
    Shortcut method to compute the KL-Divergence bewteen a model and the
    Haar distribution. The basic steps are:
        - Sample the state fidelities for randomly initialised parameters.
        - Calculates the KL divergence between the sampled probability and
          the Haar probability distribution.

    Args:
        model (Model): Function that models the quantum circuit.
        n_samples (int): Number of parameter sets to generate.
        n_bins (int): Number of histogram bins.
        random_key (Optional[jax.random.PRNGKey]): JAX random key for
            parameter initialization. If None, uses the model's internal
            random key.
        scale (bool): Whether to scale the number of samples and bins.
        kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: Tuple containing the
            input samples, bin edges, and histogram values.
    """
    _, fidelities = Expressibility.state_fidelities(
        model=model,
        random_key=random_key,
        n_samples=n_samples,
        n_bins=n_bins,
        scale=scale,
        **kwargs,
    )
    _, haar_probs = Expressibility.haar_integral(
        model.n_qubits, n_bins=n_bins, scale=scale
    )
    return Expressibility.kullback_leibler_divergence(fidelities, haar_probs)

`kullback_leibler_divergence(vqc_prob_dist, haar_dist)` `classmethod` #

Calculates the KL divergence between two probability distributions (Haar probability distribution and the fidelity distribution sampled from a VQC).

Parameters:

Name	Type	Description	Default
`vqc_prob_dist`	`ndarray`	VQC fidelity probability distribution. Should have shape (n_inputs_samples, n_bins)	required
`haar_dist`	`ndarray`	Haar probability distribution with shape. Should have shape (n_bins, )	required

Returns:

Type	Description
`ndarray`	jnp.ndarray: Array of KL-Divergence values for all values in axis 1

Source code in qml_essentials/expressibility.py

@classmethod
def kullback_leibler_divergence(
    cls,
    vqc_prob_dist: jnp.ndarray,
    haar_dist: jnp.ndarray,
) -> jnp.ndarray:
    """
    Calculates the KL divergence between two probability distributions (Haar
    probability distribution and the fidelity distribution sampled from a VQC).

    Args:
        vqc_prob_dist (jnp.ndarray): VQC fidelity probability distribution.
            Should have shape (n_inputs_samples, n_bins)
        haar_dist (jnp.ndarray): Haar probability distribution with shape.
            Should have shape (n_bins, )

    Returns:
        jnp.ndarray: Array of KL-Divergence values for all values in axis 1
    """
    if len(vqc_prob_dist.shape) > 1:
        assert all([haar_dist.shape == p.shape for p in vqc_prob_dist]), (
            "All probabilities for inputs should have the same shape as Haar. "
            f"Got {haar_dist.shape} for Haar and {vqc_prob_dist.shape} for VQC"
        )
    else:
        vqc_prob_dist = vqc_prob_dist.reshape((1, -1))

    kl_divergence = np.zeros(vqc_prob_dist.shape[0])
    for idx, p in enumerate(vqc_prob_dist):
        kl_divergence[idx] = jnp.sum(rel_entr(p, haar_dist))

    return kl_divergence

`state_fidelities(n_samples, n_bins, model, random_key=None, scale=False, **kwargs)` `classmethod` #

Sample the state fidelities and histogram them into a 2D array.

Parameters:

Name	Type	Description	Default
`n_samples`	`int`	Number of parameter sets to generate.	required
`n_bins`	`int`	Number of histogram bins.	required
`model`	`Callable`	Function that models the quantum circuit.	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`scale`	`bool`	Whether to scale the number of samples and bins.	`False`
`kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Type	Description
`ndarray`	Tuple[jnp.ndarray, jnp.ndarray]: Tuple containing the bin edges,
`ndarray`	and histogram values.

Source code in qml_essentials/expressibility.py

@classmethod
def state_fidelities(
    cls,
    n_samples: int,
    n_bins: int,
    model: Model,
    random_key: Optional[jax.random.PRNGKey] = None,
    scale: bool = False,
    **kwargs: Any,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """
    Sample the state fidelities and histogram them into a 2D array.

    Args:
        n_samples (int): Number of parameter sets to generate.
        n_bins (int): Number of histogram bins.
        model (Callable): Function that models the quantum circuit.
        random_key (Optional[jax.random.PRNGKey]): JAX random key for
            parameter initialization. If None, uses the model's internal
            random key.
        scale (bool): Whether to scale the number of samples and bins.
        kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        Tuple[jnp.ndarray, jnp.ndarray]: Tuple containing the bin edges,
        and histogram values.
    """
    if scale:
        n_samples = jnp.power(2, model.n_qubits) * n_samples
        n_bins = model.n_qubits * n_bins

    fidelities = cls._sample_state_fidelities(
        n_samples=n_samples,
        random_key=random_key,
        model=model,
        kwargs=kwargs,
    )

    y: jnp.ndarray = jnp.linspace(0, 1, n_bins + 1)

    z, _ = jnp.histogram(fidelities, bins=y)

    z = z / n_samples

    return y, z

Coefficients#

from qml_essentials.coefficients import Coefficients

Source code in qml_essentials/coefficients.py

class Coefficients:
    @classmethod
    def get_spectrum(
        cls,
        model: Model,
        mfs: int = 1,
        mts: int = 1,
        shift=False,
        trim=False,
        numerical_cap: Optional[float] = -1,
        **kwargs,
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """
        Extracts the coefficients of a given model using a FFT (jnp-fft).

        Note that the coefficients are complex numbers, but the imaginary part
        of the coefficients should be very close to zero, since the expectation
        values of the Pauli operators are real numbers.

        It can perform oversampling in both the frequency and time domain
        using the `mfs` and `mts` arguments.

        Args:
            model (Model): The model to sample.
            mfs (int): Multiplicator for the highest frequency. Default is 1.
            mts (int): Multiplicator for the number of time samples. Default is 1.
            shift (bool): Whether to apply jnp-fftshift. Default is False.
            trim (bool): Whether to remove the Nyquist frequency if spectrum is even.
                Default is False.
            numerical_cap (Optional[float]): Numerical cap for the coefficients.
                If positive, coefficients with magnitude below the cap are
                zeroed and, for a single input feature, frequencies that
                vanish entirely are removed from both `coeffs` and `freqs`.
            kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            Tuple[jnp.ndarray, jnp.ndarray]: Tuple containing the coefficients
            and frequencies.

        Note:
            The FFT grid is built from the nominal `model.degree`, which does
            not account for frequency scaling via `enc_params` or
            `enc_pulse_params`. For a model whose effective frequencies exceed
            the nominal degree by a factor $s > 1$, choose `mfs` at least
            $\\lceil s \\rceil$ (e.g. `mfs=2` covers scalings up to 2),
            otherwise the scaled components alias.
        """
        kwargs.setdefault("force_mean", True)
        kwargs.setdefault("execution_type", "expval")

        coeffs, freqs = cls._fourier_transform(model, mfs=mfs, mts=mts, **kwargs)

        if not jnp.isclose(jnp.sum(coeffs).imag, 0.0, atol=1.0e-6):
            raise ValueError(
                f"Spectrum is not real. Imaginary part of coefficients is:\
                {jnp.sum(coeffs).imag}"
            )

        if trim:
            for ax in range(model.n_input_feat):
                if coeffs.shape[ax] % 2 == 0:
                    coeffs = np.delete(coeffs, len(coeffs) // 2, axis=ax)
                    freqs = [np.delete(freq, len(freq) // 2, axis=ax) for freq in freqs]

        if shift:
            coeffs = jnp.fft.fftshift(coeffs, axes=list(range(model.n_input_feat)))
            freqs = np.fft.fftshift(freqs)

        if numerical_cap > 0:
            # set coeffs below threshold to zero
            coeffs = jnp.where(
                jnp.abs(coeffs) < numerical_cap,
                jnp.zeros_like(coeffs),
                coeffs,
            )

            # Drop frequencies whose coefficients vanish entirely after
            # capping, so the returned spectrum reflects only the surviving
            # frequencies. Well-defined only for a single (1-D) frequency
            # axis; for multi-dim input the rectangular grid is left intact.
            if model.n_input_feat == 1:
                if coeffs.ndim == 1:
                    surviving = coeffs != 0
                else:
                    surviving = jnp.any(coeffs != 0, axis=tuple(range(1, coeffs.ndim)))
                coeffs = coeffs[surviving]
                freqs = [freqs[0][surviving]]

        if len(freqs) == 1:
            freqs = freqs[0]

        return coeffs, freqs

    @classmethod
    def _fourier_transform(
        cls, model: Model, mfs: int, mts: int, **kwargs: Any
    ) -> jnp.ndarray:
        # Create a frequency vector with as many frequencies as model degrees,
        # oversampled by mfs
        n_freqs: jnp.ndarray = jnp.array(
            [mfs * model.degree[i] for i in range(model.n_input_feat)]
        )

        start, stop, step = 0, 2 * mts * jnp.pi, 2 * jnp.pi / n_freqs
        # Stretch according to the number of frequencies
        inputs: List = [
            jnp.arange(start, stop, step[i]) for i in range(model.n_input_feat)
        ]

        # permute with input dimensionality
        nd_inputs = jnp.array(
            jnp.meshgrid(*[inputs[i] for i in range(model.n_input_feat)])
        ).T.reshape(-1, model.n_input_feat)

        # Output vector is not necessarily the same length as input
        outputs = model(inputs=nd_inputs, **kwargs)
        outputs = outputs.reshape(
            *[inputs[i].shape[0] for i in range(model.n_input_feat)], -1
        ).squeeze()

        coeffs = jnp.fft.fftn(outputs, axes=list(range(model.n_input_feat)))

        freqs = [
            jnp.fft.fftfreq(int(mts * n_freqs[i]), 1 / n_freqs[i])
            for i in range(model.n_input_feat)
        ]
        # freqs = jnp.fft.fftfreq(mts * n_freqs, 1 / n_freqs)

        # TODO: this could cause issues with multidim input
        # FIXME: account for different frequencies in multidim input scenarios
        # Run the fft and rearrange +
        # normalize the output (using product if multidim)
        return (
            coeffs / math.prod(outputs.shape[0 : model.n_input_feat]),
            freqs,
        )

    @classmethod
    def get_psd(cls, coeffs: jnp.ndarray) -> jnp.ndarray:
        """
        Calculates the power spectral density (PSD) from given Fourier coefficients.

        Args:
            coeffs (jnp.ndarray): The Fourier coefficients.

        Returns:
            jnp.ndarray: The power spectral density.
        """
        # TODO: if we apply trim=True in advance, this will be slightly wrong..

        def abs2(x):
            return x.real**2 + x.imag**2

        scale = 2.0 / (len(coeffs) ** 2)
        return scale * abs2(coeffs)

    @classmethod
    def evaluate_Fourier_series(
        cls,
        coefficients: jnp.ndarray,
        frequencies: jnp.ndarray,
        inputs: Union[jnp.ndarray, list, float],
    ) -> float:
        """
        Evaluate the function value of a Fourier series at one point.

        Args:
            coefficients (jnp.ndarray): Coefficients of the Fourier series.
            frequencies (jnp.ndarray): Corresponding frequencies.
            inputs (jnp.ndarray): Point at which to evaluate the function.
        Returns:
            float: The function value at the input point.
        """
        coefficients = jnp.asarray(coefficients)

        def flatten_grid(freq_axes):
            freq_axes = [jnp.asarray(freq) for freq in freq_axes]
            freq_grid = jnp.stack(jnp.meshgrid(*freq_axes, indexing="ij"), axis=-1)
            flat_frequencies = freq_grid.reshape(-1, len(freq_axes))
            flat_coefficients = coefficients.reshape(
                flat_frequencies.shape[0], *coefficients.shape[len(freq_axes) :]
            )
            return flat_coefficients, flat_frequencies

        if isinstance(frequencies, list):
            flat_coefficients, flat_frequencies = flatten_grid(frequencies)
        else:
            frequencies = jnp.asarray(frequencies)
            if frequencies.ndim == 1:
                flat_frequencies = frequencies[:, jnp.newaxis]
                flat_coefficients = coefficients.reshape(
                    flat_frequencies.shape[0], *coefficients.shape[1:]
                )
            else:
                n_features, n_axis_freqs = frequencies.shape
                is_axis_frequencies = (
                    coefficients.shape[:n_features] == (n_axis_freqs,) * n_features
                )

                if is_axis_frequencies:
                    flat_coefficients, flat_frequencies = flatten_grid(frequencies)
                else:
                    flat_frequencies = frequencies
                    flat_coefficients = coefficients.reshape(
                        flat_frequencies.shape[0], *coefficients.shape[1:]
                    )

        inputs = jnp.asarray(inputs)
        if inputs.ndim == 0:
            inputs = inputs.reshape(1, 1)
        elif inputs.ndim == 1:
            if flat_frequencies.shape[1] == 1:
                inputs = inputs[:, jnp.newaxis]
            elif inputs.shape[0] == flat_frequencies.shape[1]:
                inputs = inputs[jnp.newaxis, :]
            else:
                inputs = jnp.repeat(
                    inputs[:, jnp.newaxis], flat_frequencies.shape[1], axis=1
                )
        exponents = jnp.exp(1j * (inputs @ flat_frequencies.T))
        exp = jnp.tensordot(exponents, flat_coefficients, axes=([1], [0]))

        return jnp.squeeze(jnp.real(exp))

`evaluate_Fourier_series(coefficients, frequencies, inputs)` `classmethod` #

Evaluate the function value of a Fourier series at one point.

Parameters:

Name	Type	Description	Default
`coefficients`	`ndarray`	Coefficients of the Fourier series.	required
`frequencies`	`ndarray`	Corresponding frequencies.	required
`inputs`	`ndarray`	Point at which to evaluate the function.	required

Returns: float: The function value at the input point.

Source code in qml_essentials/coefficients.py

@classmethod
def evaluate_Fourier_series(
    cls,
    coefficients: jnp.ndarray,
    frequencies: jnp.ndarray,
    inputs: Union[jnp.ndarray, list, float],
) -> float:
    """
    Evaluate the function value of a Fourier series at one point.

    Args:
        coefficients (jnp.ndarray): Coefficients of the Fourier series.
        frequencies (jnp.ndarray): Corresponding frequencies.
        inputs (jnp.ndarray): Point at which to evaluate the function.
    Returns:
        float: The function value at the input point.
    """
    coefficients = jnp.asarray(coefficients)

    def flatten_grid(freq_axes):
        freq_axes = [jnp.asarray(freq) for freq in freq_axes]
        freq_grid = jnp.stack(jnp.meshgrid(*freq_axes, indexing="ij"), axis=-1)
        flat_frequencies = freq_grid.reshape(-1, len(freq_axes))
        flat_coefficients = coefficients.reshape(
            flat_frequencies.shape[0], *coefficients.shape[len(freq_axes) :]
        )
        return flat_coefficients, flat_frequencies

    if isinstance(frequencies, list):
        flat_coefficients, flat_frequencies = flatten_grid(frequencies)
    else:
        frequencies = jnp.asarray(frequencies)
        if frequencies.ndim == 1:
            flat_frequencies = frequencies[:, jnp.newaxis]
            flat_coefficients = coefficients.reshape(
                flat_frequencies.shape[0], *coefficients.shape[1:]
            )
        else:
            n_features, n_axis_freqs = frequencies.shape
            is_axis_frequencies = (
                coefficients.shape[:n_features] == (n_axis_freqs,) * n_features
            )

            if is_axis_frequencies:
                flat_coefficients, flat_frequencies = flatten_grid(frequencies)
            else:
                flat_frequencies = frequencies
                flat_coefficients = coefficients.reshape(
                    flat_frequencies.shape[0], *coefficients.shape[1:]
                )

    inputs = jnp.asarray(inputs)
    if inputs.ndim == 0:
        inputs = inputs.reshape(1, 1)
    elif inputs.ndim == 1:
        if flat_frequencies.shape[1] == 1:
            inputs = inputs[:, jnp.newaxis]
        elif inputs.shape[0] == flat_frequencies.shape[1]:
            inputs = inputs[jnp.newaxis, :]
        else:
            inputs = jnp.repeat(
                inputs[:, jnp.newaxis], flat_frequencies.shape[1], axis=1
            )
    exponents = jnp.exp(1j * (inputs @ flat_frequencies.T))
    exp = jnp.tensordot(exponents, flat_coefficients, axes=([1], [0]))

    return jnp.squeeze(jnp.real(exp))

`get_psd(coeffs)` `classmethod` #

Calculates the power spectral density (PSD) from given Fourier coefficients.

Parameters:

Name	Type	Description	Default
`coeffs`	`ndarray`	The Fourier coefficients.	required

Returns:

Type	Description
`ndarray`	jnp.ndarray: The power spectral density.

Source code in qml_essentials/coefficients.py

@classmethod
def get_psd(cls, coeffs: jnp.ndarray) -> jnp.ndarray:
    """
    Calculates the power spectral density (PSD) from given Fourier coefficients.

    Args:
        coeffs (jnp.ndarray): The Fourier coefficients.

    Returns:
        jnp.ndarray: The power spectral density.
    """
    # TODO: if we apply trim=True in advance, this will be slightly wrong..

    def abs2(x):
        return x.real**2 + x.imag**2

    scale = 2.0 / (len(coeffs) ** 2)
    return scale * abs2(coeffs)

`get_spectrum(model, mfs=1, mts=1, shift=False, trim=False, numerical_cap=-1, **kwargs)` `classmethod` #

Extracts the coefficients of a given model using a FFT (jnp-fft).

Note that the coefficients are complex numbers, but the imaginary part of the coefficients should be very close to zero, since the expectation values of the Pauli operators are real numbers.

It can perform oversampling in both the frequency and time domain using the mfs and mts arguments.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The model to sample.	required
`mfs`	`int`	Multiplicator for the highest frequency. Default is 1.	`1`
`mts`	`int`	Multiplicator for the number of time samples. Default is 1.	`1`
`shift`	`bool`	Whether to apply jnp-fftshift. Default is False.	`False`
`trim`	`bool`	Whether to remove the Nyquist frequency if spectrum is even. Default is False.	`False`
`numerical_cap`	`Optional[float]`	Numerical cap for the coefficients. If positive, coefficients with magnitude below the cap are zeroed and, for a single input feature, frequencies that vanish entirely are removed from both `coeffs` and `freqs`.	`-1`
`kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Type	Description
`ndarray`	Tuple[jnp.ndarray, jnp.ndarray]: Tuple containing the coefficients
`ndarray`	and frequencies.

Note

The FFT grid is built from the nominal model.degree, which does not account for frequency scaling via enc_params or enc_pulse_params. For a model whose effective frequencies exceed the nominal degree by a factor \(s > 1\), choose mfs at least \(\lceil s \rceil\) (e.g. mfs=2 covers scalings up to 2), otherwise the scaled components alias.

Source code in qml_essentials/coefficients.py

@classmethod
def get_spectrum(
    cls,
    model: Model,
    mfs: int = 1,
    mts: int = 1,
    shift=False,
    trim=False,
    numerical_cap: Optional[float] = -1,
    **kwargs,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """
    Extracts the coefficients of a given model using a FFT (jnp-fft).

    Note that the coefficients are complex numbers, but the imaginary part
    of the coefficients should be very close to zero, since the expectation
    values of the Pauli operators are real numbers.

    It can perform oversampling in both the frequency and time domain
    using the `mfs` and `mts` arguments.

    Args:
        model (Model): The model to sample.
        mfs (int): Multiplicator for the highest frequency. Default is 1.
        mts (int): Multiplicator for the number of time samples. Default is 1.
        shift (bool): Whether to apply jnp-fftshift. Default is False.
        trim (bool): Whether to remove the Nyquist frequency if spectrum is even.
            Default is False.
        numerical_cap (Optional[float]): Numerical cap for the coefficients.
            If positive, coefficients with magnitude below the cap are
            zeroed and, for a single input feature, frequencies that
            vanish entirely are removed from both `coeffs` and `freqs`.
        kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        Tuple[jnp.ndarray, jnp.ndarray]: Tuple containing the coefficients
        and frequencies.

    Note:
        The FFT grid is built from the nominal `model.degree`, which does
        not account for frequency scaling via `enc_params` or
        `enc_pulse_params`. For a model whose effective frequencies exceed
        the nominal degree by a factor $s > 1$, choose `mfs` at least
        $\\lceil s \\rceil$ (e.g. `mfs=2` covers scalings up to 2),
        otherwise the scaled components alias.
    """
    kwargs.setdefault("force_mean", True)
    kwargs.setdefault("execution_type", "expval")

    coeffs, freqs = cls._fourier_transform(model, mfs=mfs, mts=mts, **kwargs)

    if not jnp.isclose(jnp.sum(coeffs).imag, 0.0, atol=1.0e-6):
        raise ValueError(
            f"Spectrum is not real. Imaginary part of coefficients is:\
            {jnp.sum(coeffs).imag}"
        )

    if trim:
        for ax in range(model.n_input_feat):
            if coeffs.shape[ax] % 2 == 0:
                coeffs = np.delete(coeffs, len(coeffs) // 2, axis=ax)
                freqs = [np.delete(freq, len(freq) // 2, axis=ax) for freq in freqs]

    if shift:
        coeffs = jnp.fft.fftshift(coeffs, axes=list(range(model.n_input_feat)))
        freqs = np.fft.fftshift(freqs)

    if numerical_cap > 0:
        # set coeffs below threshold to zero
        coeffs = jnp.where(
            jnp.abs(coeffs) < numerical_cap,
            jnp.zeros_like(coeffs),
            coeffs,
        )

        # Drop frequencies whose coefficients vanish entirely after
        # capping, so the returned spectrum reflects only the surviving
        # frequencies. Well-defined only for a single (1-D) frequency
        # axis; for multi-dim input the rectangular grid is left intact.
        if model.n_input_feat == 1:
            if coeffs.ndim == 1:
                surviving = coeffs != 0
            else:
                surviving = jnp.any(coeffs != 0, axis=tuple(range(1, coeffs.ndim)))
            coeffs = coeffs[surviving]
            freqs = [freqs[0][surviving]]

    if len(freqs) == 1:
        freqs = freqs[0]

    return coeffs, freqs

Fourier Tree#

from qml_essentials.coefficients import FourierTree

Sine-cosine tree representation for the algorithm by Nemkov et al.

Computes the analytical Fourier coefficients/frequencies of a Pauli-Clifford circuit. The symbolic structure of the tree (which Pauli rotations contribute sine/cosine factors to which leaf, and the leaf observables) is built once in NumPy; the parameter-dependent coefficients are then obtained with a small number of vectorised JAX operations, so the result remains jittable / differentiable with respect to the model parameters.

The resulting spectrum is the d-dimensional set of frequency vectors, where \(d\) is the input dimensionality.

Usage:

model = Model(...)
tree = FourierTree(model)
exp = tree()                          # expectation value
coeff_list, freq_list = tree.get_spectrum()

Source code in qml_essentials/coefficients.py

class FourierTree:
    """
    Sine-cosine tree representation for the algorithm by Nemkov et al.

    Computes the analytical Fourier coefficients/frequencies of a Pauli-Clifford
    circuit.  The symbolic structure of the tree (which Pauli rotations
    contribute sine/cosine factors to which leaf, and the leaf observables) is
    built once in NumPy; the parameter-dependent coefficients are then obtained
    with a small number of vectorised JAX operations, so the result remains
    jittable / differentiable with respect to the model parameters.

    The resulting spectrum is the d-dimensional set of frequency vectors,
    where $d$ is the input dimensionality.

    **Usage**:
    ```
    model = Model(...)
    tree = FourierTree(model)
    exp = tree()                          # expectation value
    coeff_list, freq_list = tree.get_spectrum()
    ```
    """

    def __init__(self, model: Model):
        """
        Tree initialisation, based on the Pauli-Clifford representation of a
        model.

        Args:
            model (Model): The Model, for which to build the tree.
        """
        self.model = model
        self.n_qubits = model.n_qubits

        # A single (de-batched) parameter set drives the whole tree.
        self._params = self._single_param_set(model.params)

        # Canonical Pauli-Clifford structure, recorded once at a fixed base
        # input.  The base value is irrelevant to the structure (it only sets
        # the rotation angles, not which Pauli words appear).
        base_inputs = np.ones(model.n_input_feat)
        operations, observables = self._build_canonical_tape(self._params, base_inputs)

        self.parameters = [
            jnp.squeeze(p) for p in PauliCircuit.get_parameters(operations)
        ]
        self.n_params = len(self.parameters)

        # Pauli generators of the (canonical) rotations, as symbolic words.
        self.pauli_words: List[PauliWord] = [
            PauliWord.from_operation(op, self.n_qubits) for op in operations
        ]

        # Cumulative X/Y support of the rotations[0..k] (for light-cone early
        # stopping).  cumulative_xy[k] is True on every qubit touched by an X/Y
        # generator in any rotation up to index k.
        self.cumulative_xy: List[np.ndarray] = []
        running = np.zeros(self.n_qubits, dtype=bool)
        for pw in self.pauli_words:
            running = np.logical_or(running, pw.xy_mask)
            self.cumulative_xy.append(running.copy())

        # Observable Pauli words (one tree root each).
        self.observable_words: List[PauliWord] = [
            PauliWord.from_operation(obs, self.n_qubits) for obs in observables
        ]

        # Identify the input-encoding columns, their feature, and integer
        # frequency scaling directly from the tape (no per-gate tagging).  Sets
        # ``input_indices``, ``all_input_indices``, ``input_scaling``,
        # ``var_positions`` and ``features``.
        self._detect_inputs(base_inputs)

        # The explicit leaf structure is built lazily: for deep circuits the
        # number of tree paths explodes combinatorially, while the canonical
        # form above (and the merged-state support DP) remain cheap.
        self._structure_built = False

    def _ensure_structure(self) -> None:
        """Build the explicit leaf/spectrum structure on first use."""
        if not self._structure_built:
            # Symbolic structure: per root (S, C, terms) leaf arrays ...
            self._build_leaf_arrays()
            # ... and the parameter-independent frequency/weight structure.
            self._build_spectrum_structure()
            self._structure_built = True

    def _single_param_set(self, params) -> jnp.ndarray:
        """De-batch the model parameters to the single set the tree describes.

        Models can carry batched parameters (e.g. after FCC sampling); the tree
        is defined for one set, so fall back to the first and warn.
        """
        params = jnp.asarray(params)
        if params.ndim > 2 and params.shape[0] > 1:
            warnings.warn(
                f"FourierTree supports a single parameter set; using the first "
                f"of {params.shape[0]} batched parameter sets.",
                UserWarning,
            )
            params = params[0]
        return params

    def _build_canonical_tape(self, params, inputs):
        """Record the circuit and transform it to Pauli-Clifford normal form.

        Returns the ``(operations, observables)`` of the canonical circuit
        (see :meth:`PauliCircuit.from_parameterised_circuit`).
        """
        params = self._single_param_set(params)
        inputs = self.model._inputs_validation(inputs)
        raw_tape = self.model.script.record(params=params, inputs=inputs)
        _, obs_list = self.model._build_obs()
        return PauliCircuit.from_parameterised_circuit(
            raw_tape, observables=obs_list, n_qubits=self.n_qubits
        )

    def _canonical_parameters(self, inputs) -> np.ndarray:
        """Recorded canonical rotation angles (1-D float array) for ``inputs``."""
        operations, _ = self._build_canonical_tape(self._params, inputs)
        return np.array(
            [float(jnp.squeeze(p)) for p in PauliCircuit.get_parameters(operations)]
        )

    def _detect_inputs(self, base_inputs: np.ndarray) -> None:
        r"""Infer the input-encoding columns directly from the tape (tag-free).

        Each encoding rotation applies an angle :math:`\omega_k\,x_f` that is
        linear in a single input feature :math:`x_f`, and Clifford commutation
        only multiplies a rotation generator by :math:`\pm 1`.  Every canonical
        rotation angle is therefore an affine function of the inputs, so
        perturbing one feature at a time and differencing the recorded angles
        isolates exactly the columns that depend on it, together with the
        signed scaling :math:`\omega_k` (rational for diagonal-Hamiltonian
        encodings such as Golomb).

        Sets :attr:`input_indices` (``{feature: [columns]}``),
        :attr:`all_input_indices`, :attr:`input_scaling` (per column, ``1`` for
        variational columns), :attr:`var_positions`, and :attr:`features`.

        Raises:
            NotImplementedError: If a rotation depends on more than one feature
                (the tree requires single-feature encodings).
        """
        tol = 1e-6
        d = self.model.n_input_feat
        base = np.asarray(base_inputs, dtype=float)
        p_base = np.array([float(p) for p in self.parameters])

        # response[f, k] = d(angle_k) / d(x_f), the linear response of column k.
        response = np.zeros((d, self.n_params))
        for f in range(d):
            step = base.copy()
            step[f] += 1.0
            response[f] = self._canonical_parameters(step) - p_base

        input_indices: Dict[int, list] = defaultdict(list)
        all_input_indices: List[int] = []
        scaling = np.ones(self.n_params, dtype=np.float64)
        for k in range(self.n_params):
            feats = np.flatnonzero(np.abs(response[:, k]) > tol)
            if feats.size == 0:
                continue  # variational column
            if feats.size > 1:
                raise NotImplementedError(
                    f"Rotation {k} depends on multiple input features "
                    f"{feats.tolist()}; the Fourier tree requires each encoding "
                    "rotation to be linear in a single feature."
                )
            f = int(feats[0])
            omega = float(response[f, k])
            # Scalings may be rational (e.g. diagonal-Hamiltonian / Golomb
            # encodings, whose per-Pauli-string scalings are dyadic rationals);
            # the integer-frequency model spectrum is recovered downstream.
            input_indices[f].append(k)
            all_input_indices.append(k)
            scaling[k] = omega

        self.input_indices = input_indices
        self.all_input_indices = all_input_indices
        self.input_scaling = scaling
        input_set = set(all_input_indices)
        self.var_positions = np.array(
            [i for i in range(self.n_params) if i not in input_set], dtype=np.int64
        )
        # Ordered list of input feature keys (d-dimensional spectrum).
        self.features = sorted(input_indices.keys())

    # Symbolic tree construction (NumPy)
    def _build_leaf_arrays(self) -> None:
        """Collect the tree leaves for every root into integer count matrices.

        For each root (observable) this produces:
            - ``S``: (n_leaves, n_params) sine-factor counts per parameter,
            - ``C``: (n_leaves, n_params) cosine-factor counts per parameter,
            - ``terms``: (n_leaves,) complex leaf constants ``<0|O_leaf|0>``.
        """
        self.leaf_arrays: List[Tuple[np.ndarray, np.ndarray, np.ndarray]] = []
        for obs_word in self.observable_words:
            leaves: List[Tuple[np.ndarray, np.ndarray, complex]] = []
            zeros = np.zeros(self.n_params, dtype=np.int64)
            self._collect_leaves(
                obs_word, self.n_params - 1, zeros.copy(), zeros.copy(), leaves
            )
            if leaves:
                S = np.stack([leaf[0] for leaf in leaves])
                C = np.stack([leaf[1] for leaf in leaves])
                terms = np.array([leaf[2] for leaf in leaves], dtype=np.complex128)
            else:
                S = np.zeros((0, self.n_params), dtype=np.int64)
                C = np.zeros((0, self.n_params), dtype=np.int64)
                terms = np.zeros(0, dtype=np.complex128)
            self.leaf_arrays.append((S, C, terms))

    def _collect_leaves(
        self,
        observable: PauliWord,
        pauli_idx: int,
        sin_counts: np.ndarray,
        cos_counts: np.ndarray,
        leaves: List[Tuple[np.ndarray, np.ndarray, complex]],
    ) -> None:
        """Recursively enumerate the leaves of the coefficient tree.

        The incoming sine/cosine factor (from the parent edge) is already
        accumulated into ``sin_counts``/``cos_counts``.  This fuses the tree
        construction and leaf traversal of the original implementation into a
        single NumPy pass (no per-node JAX scatter updates).
        """
        if self._early_stopping_possible(pauli_idx, observable):
            return

        # Skip trailing Pauli rotations that commute with the observable.
        while pauli_idx >= 0:
            last = self.pauli_words[pauli_idx]
            if not observable.commutes_with(last):
                break
            pauli_idx -= 1
        else:  # leaf reached
            term = observable.zero_expectation()
            if term != 0:
                leaves.append((sin_counts, cos_counts, term))
            return

        last = self.pauli_words[pauli_idx]

        # Left child: cosine factor for this parameter, same observable.
        cos_left = cos_counts.copy()
        cos_left[pauli_idx] += 1
        self._collect_leaves(
            observable, pauli_idx - 1, sin_counts.copy(), cos_left, leaves
        )

        # Right child: sine factor, observable becomes  P . O.
        sin_right = sin_counts.copy()
        sin_right[pauli_idx] += 1
        self._collect_leaves(
            last.compose(observable),
            pauli_idx - 1,
            sin_right,
            cos_counts.copy(),
            leaves,
        )

    def _early_stopping_possible(self, pauli_idx: int, observable: PauliWord) -> bool:
        """Whether a node can be discarded (all reachable expectations vanish).

        Mirrors the criterion of Nemkov et al. (light cone): a qubit on which
        the observable carries an X/Y must be covered by an X/Y generator of
        some remaining rotation (rotations[0..pauli_idx]); otherwise that X/Y can
        never be rotated into a diagonal term and the whole node contributes
        zero.  Equivalently, the node survives iff every qubit is either I/Z in
        the observable or covered by the cumulative rotation X/Y support.
        """
        obs_iz = np.logical_not(observable.xy_mask)
        combined = np.logical_or(obs_iz, self.cumulative_xy[pauli_idx]).all()
        return not bool(combined)

    # Frequency / weight structure (NumPy, parameter independent)
    def _build_spectrum_structure(self) -> None:
        """Build, per root, the frequency vectors and the (n_freq, n_leaves)
        weight matrix ``W`` such that ``coeffs = W @ (terms * variational)``.
        """
        self.freqs_per_root: List[np.ndarray] = []
        self.weights_per_root: List[np.ndarray] = []
        d = len(self.features)

        for S, C, _ in self.leaf_arrays:
            n_leaves = S.shape[0]
            freq_to_col: Dict[tuple, np.ndarray] = defaultdict(
                lambda: np.zeros(n_leaves, dtype=np.complex128)
            )
            for leaf in range(n_leaves):
                # One expansion factor per *active* input column, each carrying
                # its feature axis and integer frequency scaling.  Per leaf a
                # column contributes at most one sin/cos factor (square-free),
                # but different columns of the same feature may carry different
                # scalings, so they are expanded individually and convolved
                # rather than aggregating counts (which would assume a common
                # unit scaling).
                col_factors: List[List[Tuple[int, float, float]]] = []
                half_exp = 0
                for axis, feat in enumerate(self.features):
                    for k in self.input_indices[feat]:
                        s = int(S[leaf, k])
                        c = int(C[leaf, k])
                        if s == 0 and c == 0:
                            continue
                        half_exp += s + c
                        w_k = float(self.input_scaling[k])
                        col_factors.append(
                            [
                                (axis, o * w_k, wt)
                                for o, wt in self._binomial_terms(s, c)
                            ]
                        )
                half = 0.5**half_exp

                if d == 0:
                    freq_to_col[(0,)][leaf] += half
                    continue

                if not col_factors:
                    freq_to_col[(0,) * d][leaf] += half
                    continue

                for combo in itertools.product(*col_factors):
                    omega = [0.0] * d
                    weight = half
                    for axis, o, wt in combo:
                        omega[axis] += o
                        weight *= wt
                    # Snap to a tolerance so dyadic-rational contributions that
                    # are numerically equal share a single frequency key.
                    key = tuple(round(v, 9) for v in omega)
                    freq_to_col[key][leaf] += weight

            if freq_to_col:
                omegas = sorted(freq_to_col.keys())
                W = np.stack([freq_to_col[o] for o in omegas])  # (n_freq, n_leaves)
                freqs = np.array(omegas, dtype=float)  # (n_freq, d)
                # Diagonal-Hamiltonian (Golomb) encodings produce rational
                # per-rotation scalings but integer model frequencies; snap to
                # int64 when every entry is integral so integer-encoding
                # consumers are unchanged, otherwise keep the rational floats.
                rounded = np.rint(freqs)
                if np.all(np.abs(freqs - rounded) < 1e-6):
                    freqs = rounded.astype(np.int64)
            else:
                freqs = np.zeros((1, max(d, 1)), dtype=np.int64)
                W = np.zeros((1, n_leaves), dtype=np.complex128)

            # Collapse to 1-D frequency array for the single-feature case.
            if freqs.shape[1] == 1:
                freqs = freqs[:, 0]
            self.freqs_per_root.append(freqs)
            # Keep W in NumPy complex128: its entries are dyadic rationals
            # (binomial weights x 0.5^k x i^m), which are exact in float64 --
            # this allows exact symbolic zero-tests in get_exact_support.
            self.weights_per_root.append(W)

    @staticmethod
    def _binomial_terms(s: int, c: int) -> List[Tuple[int, float]]:
        """Expand ``cos^c (i sin)^s`` in ``e^{i omega x}`` (without the 0.5 factor).

        Returns a list of ``(omega, weight)`` with
        ``omega = 2a + 2b - s - c`` and ``weight = C(s,a) C(c,b) (-1)^{s-a}``.
        """
        terms = []
        for a in range(s + 1):
            for b in range(c + 1):
                weight = math.comb(s, a) * math.comb(c, b) * (-1) ** (s - a)
                terms.append((2 * a + 2 * b - s - c, float(weight)))
        return terms

    # Vectorised numeric evaluation (JAX)
    @staticmethod
    def _safe_pow(base: jnp.ndarray, exp: jnp.ndarray) -> jnp.ndarray:
        """Elementwise ``base ** exp`` for real base and non-negative integer
        exponents, correct for negative bases (avoids ``log`` of negatives).

        Args:
            base: real array of shape ``(n,)``.
            exp: integer array of shape ``(n_leaves, n)``.
        """
        mag = jnp.abs(base)[None, :] ** exp
        sign = jnp.where(exp % 2 == 0, 1.0, jnp.sign(base)[None, :])
        return sign * mag

    _I_POW = None  # set lazily to jnp.array([1, 1j, -1, -1j])

    def _leaf_factors(
        self, S: np.ndarray, C: np.ndarray, columns: np.ndarray
    ) -> jnp.ndarray:
        """Per-leaf product ``prod_i cos(theta_i)^{C} (i sin(theta_i))^{S}`` over
        the given parameter ``columns`` (vectorised over leaves).
        """
        if FourierTree._I_POW is None:
            FourierTree._I_POW = jnp.array([1, 1j, -1, -1j])

        if S.shape[0] == 0:
            return jnp.zeros(0, dtype=jnp.complex64)

        theta = jnp.stack([self.parameters[i] for i in columns])
        S_sub = jnp.asarray(S[:, columns])
        C_sub = jnp.asarray(C[:, columns])

        cos_part = self._safe_pow(jnp.cos(theta), C_sub)
        sin_mag = self._safe_pow(jnp.sin(theta), S_sub)
        i_part = FourierTree._I_POW[S_sub % 4]
        return jnp.prod(cos_part * sin_mag * i_part, axis=1)

    def __call__(
        self,
        params: Optional[jnp.ndarray] = None,
        inputs: Optional[jnp.ndarray] = None,
        **kwargs,
    ) -> jnp.ndarray:
        """
        Evaluate the expectation value(s) of the model's observables via the
        sine-cosine tree (equivalent to the circuit expectation).

        Args:
            params (Optional[jnp.ndarray]): Model parameters. Defaults to the
                model's parameters.
            inputs (Optional[jnp.ndarray]): Inputs to the circuit. Defaults to 1.

        Returns:
            jnp.ndarray: Expectation value per observable (or their mean if
                ``force_mean`` is set).

        Raises:
            NotImplementedError: For execution types other than "expval" or when
                noise is requested.
        """
        params = (
            self.model._params_validation(params)
            if params is not None
            else self.model.params
        )
        inputs = (
            self.model._inputs_validation(inputs)
            if inputs is not None
            else self.model._inputs_validation(1.0)
        )

        if kwargs.get("execution_type", "expval") != "expval":
            raise NotImplementedError(
                f'Currently, only "expval" execution type is supported when '
                f"building FourierTree. Got {kwargs.get('execution_type', 'expval')}."
            )
        if kwargs.get("noise_params", None) is not None:
            raise NotImplementedError(
                "Currently, noise is not supported when building FourierTree."
            )

        # Re-derive the (canonical) parameter values for the requested inputs;
        # the tree structure (leaf arrays) is unchanged.
        operations, _ = self._build_canonical_tape(params, inputs)
        self.parameters = [
            jnp.squeeze(p) for p in PauliCircuit.get_parameters(operations)
        ]

        self._ensure_structure()
        all_columns = np.arange(self.n_params, dtype=np.int64)
        results = []
        for S, C, terms in self.leaf_arrays:
            factors = self._leaf_factors(S, C, all_columns)
            results.append(jnp.real(jnp.sum(jnp.asarray(terms) * factors)))
        results = jnp.array(results)

        if kwargs.get("force_mean", False):
            return jnp.mean(results)
        return results

    def get_spectrum(
        self, force_mean: bool = False
    ) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
        """
        Compute the Fourier spectrum (coefficients and frequencies) of the tree.

        Args:
            force_mean (bool, optional): Average the coefficients over all
                observables (roots). Defaults to False.

        Returns:
            Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
                - List of coefficients, one entry per observable (root).
                - List of corresponding frequencies, one entry per root.
                When ``force_mean`` is set, both lists have a single entry.
        """
        self._ensure_structure()
        per_root_coeffs: List[jnp.ndarray] = []
        for (S, C, terms), W in zip(self.leaf_arrays, self.weights_per_root):
            leaf_const = jnp.asarray(terms) * self._leaf_factors(
                S, C, self.var_positions
            )
            per_root_coeffs.append(jnp.asarray(W) @ leaf_const)

        return self._combine_roots(per_root_coeffs, self.freqs_per_root, force_mean)

    def get_exact_support(self, method: str = "tree") -> List[np.ndarray]:
        r"""Symbolically derive the exact frequency support (no sampling).

        A frequency :math:`\omega` belongs to the exact spectrum iff its
        coefficient :math:`c_\omega(\theta) = \sum_l W_{\omega l}\,
        \text{term}_l\, v_l(\theta)` is not identically zero in the
        variational parameters :math:`\theta`.

        Two methods are available:

        - ``"tree"`` (default, fully exact): enumerates the explicit tree
          leaves.  Because the branch index strictly decreases along every tree
          path, each parameter contributes **at most one** sine *or* cosine
          factor per leaf (:math:`S_{li}, C_{li} \in \{0, 1\}`).  Every
          variational leaf factor :math:`v_l` is therefore a *square-free*
          monomial over :math:`\{1, \cos\theta_i, i\sin\theta_i\}`, and
          monomials with distinct signatures are linearly independent functions
          (no :math:`\cos^2 + \sin^2` identities can arise without squares).
          Hence

          .. math::
              c_\omega \equiv 0 \iff \sum_{l \in g} W_{\omega l}\,\text{term}_l
              = 0 \quad \text{for every signature group } g.

          Since all involved quantities are dyadic rationals times
          :math:`\{\pm 1, \pm i\}`, the group sums are exact in float64 and the
          zero-test is exact.  The number of leaves can however grow
          exponentially with circuit depth.

        - ``"dp"`` (scalable): merges tree nodes with identical
          ``(rotation index, observable)`` — at most ``n_params * 4^n_qubits``
          states — and tracks, per state, the achievable per-feature sine/cosine
          count vectors ``(s_f, c_f)`` as a mixed-radix bitmask.  Each feature's
          support is the union of the (exact) expansion supports of
          :math:`\cos^{c_f} x_f\, (i \sin x_f)^{s_f}`, and the model support is
          their Cartesian product across features.  This is exact per tree path
          (including interior zero coefficients of the expansions), but unlike
          ``"tree"`` it cannot detect coefficients that cancel identically
          *across* paths with identical variational signatures (e.g. directly
          repeated encodings).  It therefore yields a tight superset in such
          corner cases.  Supports any number of input features, but requires
          unit-magnitude input scaling: per-gate :math:`|\omega| \neq 1`
          scalings (e.g. Golomb encodings) are rejected — use ``"tree"``.

        Args:
            method (str): ``"tree"`` (fully exact) or ``"dp"`` (scalable).

        Returns:
            List[np.ndarray]: For each observable (root), the frequency vectors
            with not-identically-zero coefficient — shape ``(n_freq,)`` for a
            single input feature, ``(n_freq, n_features)`` otherwise.
        """
        if method == "dp":
            return self._support_dp()
        if method != "tree":
            raise ValueError(f"Unknown method '{method}'. Use 'tree' or 'dp'.")

        self._ensure_structure()
        supports = []
        for (S, C, terms), W, freqs in zip(
            self.leaf_arrays, self.weights_per_root, self.freqs_per_root
        ):
            freqs = np.asarray(freqs)
            n_leaves = S.shape[0]
            if n_leaves == 0:
                supports.append(freqs[:0])
                continue

            # Group leaves by their variational sine/cosine signature.
            signature = np.hstack([S[:, self.var_positions], C[:, self.var_positions]])
            _, groups = np.unique(signature, axis=0, return_inverse=True)
            n_groups = int(groups.max()) + 1

            # Per-group sums of W[omega, l] * term_l, accumulated exactly.
            contrib = (W * terms[None, :]).T  # (n_leaves, n_freq)
            group_sums = np.zeros((n_groups, W.shape[0]), dtype=np.complex128)
            np.add.at(group_sums, groups, contrib)

            mask = (np.abs(group_sums) > 1e-12).any(axis=0)  # (n_freq,)
            supports.append(freqs[mask])
        return supports

    def _support_dp(self) -> List[np.ndarray]:
        """Merged-state dynamic program for the frequency support.

        Instead of enumerating all (worst-case exponentially many) tree paths,
        nodes are merged on ``(rotation index, bare observable)``.  Each state
        stores the set of achievable per-axis count vectors
        ``(c_0, s_0, ..., c_{d-1}, s_{d-1})`` as a mixed-radix bitmask (one
        digit per feature sine/cosine count), so transitions are O(1) big-int
        operations.  See :meth:`get_exact_support` for semantics and
        limitations.
        """
        # Count aggregation is valid as long as every input rotation has
        # unit-magnitude scaling.  A Clifford-commutation sign flip (scaling -1)
        # leaves the frequency support unchanged -- cos is even and sin odd, so
        # the sign only flips the coefficient, which the support ignores -- but a
        # genuine per-gate scaling (|omega| != 1, e.g. Golomb / heterogeneous
        # frequencies) cannot be represented by sin/cos counts.
        if self.all_input_indices and np.any(
            np.abs(self.input_scaling[self.all_input_indices]) != 1
        ):
            raise NotImplementedError(
                "The 'dp' support method does not support non-unit input "
                "frequency scaling (it aggregates sin/cos counts and cannot "
                "represent per-gate scalings); use method='tree'."
            )

        n = self.n_qubits
        d = len(self.features)
        # Pack the achievable per-axis sine/cosine counts into one big-int as a
        # mixed-radix bitmask over (c_0, s_0, ..., c_{d-1}, s_{d-1}).  Axis a's
        # counts range 0..n_a (n_a input rotations encode feature a), so each
        # digit has radix n_a + 1; a left-shift by a digit's place value
        # increments that count, OR is set-union.  (For d == 1 this reduces to
        # shift_c = 1, shift_s = n_inp + 1, i.e. the flat (s, c) layout.)
        ranges = [len(self.input_indices[self.features[a]]) + 1 for a in range(d)]
        shift_c = [0] * d
        shift_s = [0] * d
        place = 1
        for a in range(d):
            shift_c[a] = place
            place *= ranges[a]
            shift_s[a] = place
            place *= ranges[a]
        # Feature axis of each input rotation (-1 if variational), in the same
        # enumerate(self.features) order the tree method uses for its axes.
        axis_of_col = np.full(self.n_params, -1, dtype=np.int64)
        for a in range(d):
            for k in self.input_indices[self.features[a]]:
                axis_of_col[k] = a

        def encode(word: PauliWord) -> Tuple[int, int]:
            x = z = 0
            for q in range(n):
                x |= int(word.x[q]) << q
                z |= int(word.z[q]) << q
            return x, z

        paulis = [encode(w) for w in self.pauli_words]
        cum_xy = []
        running = 0
        for xp, _ in paulis:
            running |= xp
            cum_xy.append(running)

        def parity(v: int) -> int:
            return bin(v).count("1") & 1

        def dp(idx: int, xo: int, zo: int, memo: dict) -> int:
            # Light-cone early stopping (cf. _early_stopping_possible).
            if idx >= 0 and (xo & ~cum_xy[idx]):
                return 0
            # Skip trailing rotations that commute with the observable.
            while idx >= 0:
                xp, zp = paulis[idx]
                if parity(xo & zp) ^ parity(zo & xp):
                    break
                idx -= 1
            else:  # leaf: counts (s=0, c=0) iff the observable is diagonal
                return 1 if xo == 0 else 0
            key = (idx, xo, zo)
            hit = memo.get(key)
            if hit is not None:
                return hit
            xp, zp = paulis[idx]
            cos_child = dp(idx - 1, xo, zo, memo)
            sin_child = dp(idx - 1, xo ^ xp, zo ^ zp, memo)
            a = int(axis_of_col[idx])
            if a >= 0:
                # Active input gate: cosine increments c_a, sine increments s_a.
                val = (cos_child << shift_c[a]) | (sin_child << shift_s[a])
            else:
                val = cos_child | sin_child
            memo[key] = val
            return val

        # Recursion depth is bounded by the number of rotations.
        old_limit = sys.getrecursionlimit()
        sys.setrecursionlimit(max(old_limit, self.n_params + 1000))
        try:
            supports = []
            for obs in self.observable_words:
                memo: dict = {}
                xo, zo = encode(obs)
                mask = dp(self.n_params - 1, xo, zo, memo)
                supports.append(self._dp_mask_to_support(mask, d, ranges))
        finally:
            sys.setrecursionlimit(old_limit)
        return supports

    def _dp_mask_to_support(self, mask: int, d: int, ranges: List[int]) -> np.ndarray:
        """Decode a count bitmask (see :meth:`_support_dp`) into a frequency
        support.  Each set bit is a per-axis count vector ``(c_a, s_a)``; the
        per-axis expansion supports are combined across axes (Cartesian
        product) and unioned over all bits.

        Returns an ``(n_freq, d)`` array, collapsed to 1-D for ``d <= 1`` to
        match :meth:`get_exact_support` with ``method="tree"`` (``d == 0`` keeps
        only the DC term).
        """
        tupleset: set = set()
        while mask:
            bit = mask & -mask
            i = bit.bit_length() - 1
            rem = i
            axis_freqs = []
            for a in range(d):
                c_a = rem % ranges[a]
                rem //= ranges[a]
                s_a = rem % ranges[a]
                rem //= ranges[a]
                axis_freqs.append(sorted(self._expansion_support(s_a, c_a)))
            tupleset.update(itertools.product(*axis_freqs))
            mask ^= bit

        if d >= 2:
            if not tupleset:
                return np.empty((0, d), dtype=np.int64)
            return np.array(sorted(tupleset), dtype=np.int64)
        # d in {0, 1}: collapse to a 1-D frequency array.
        flat = sorted(t[0] for t in tupleset) if d == 1 else ([0] if tupleset else [])
        return np.array(flat, dtype=np.int64)

    @staticmethod
    @lru_cache(maxsize=None)
    def _expansion_support(s: int, c: int) -> frozenset:
        r"""Frequencies with non-zero coefficient in :math:`\cos^c x (i\sin x)^s`.

        Computed exactly with integer arithmetic via the polynomial
        :math:`(t - 1)^s (t + 1)^c` (with :math:`t = e^{2ix}` up to a shift);
        interior coefficients can vanish, e.g. :math:`\cos x \sin x` only
        contains :math:`\pm 2`.
        """
        coeffs = [1]
        for _ in range(s):  # multiply by (t - 1)
            new = [0] * (len(coeffs) + 1)
            for i, a in enumerate(coeffs):
                new[i + 1] += a
                new[i] -= a
            coeffs = new
        for _ in range(c):  # multiply by (t + 1)
            new = [0] * (len(coeffs) + 1)
            for i, a in enumerate(coeffs):
                new[i + 1] += a
                new[i] += a
            coeffs = new
        m = s + c
        return frozenset(2 * k - m for k, a in enumerate(coeffs) if a != 0)

    def _combine_roots(
        self,
        per_root_coeffs: List[jnp.ndarray],
        per_root_freqs: List[np.ndarray],
        force_mean: bool,
    ) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
        """Assemble the per-root spectra, optionally averaging over roots."""
        if not force_mean:
            coefficients = [jnp.asarray(c) for c in per_root_coeffs]
            frequencies = [jnp.asarray(f) for f in per_root_freqs]
            return coefficients, frequencies

        # Average over roots on the union of all frequency vectors.  Keys are
        # snapped to a tolerance so rational frequencies dedup robustly.
        accum: Dict[tuple, complex] = defaultdict(complex)
        for coeffs, freqs in zip(per_root_coeffs, per_root_freqs):
            freqs_np = np.asarray(freqs)
            for k in range(freqs_np.shape[0]):
                key = (
                    (round(float(freqs_np[k]), 9),)
                    if freqs_np.ndim == 1
                    else tuple(round(float(v), 9) for v in freqs_np[k])
                )
                accum[key] += complex(coeffs[k])
        n_roots = max(len(per_root_coeffs), 1)
        keys = sorted(accum.keys())
        mean_coeffs = jnp.array([accum[k] / n_roots for k in keys])
        freq_arr = np.array(keys, dtype=float)
        rounded = np.rint(freq_arr)
        if np.all(np.abs(freq_arr - rounded) < 1e-6):
            freq_arr = rounded.astype(np.int64)
        if freq_arr.shape[1] == 1:
            freq_arr = freq_arr[:, 0]
        return [mean_coeffs], [jnp.asarray(freq_arr)]

`call(params=None, inputs=None, **kwargs)` #

Evaluate the expectation value(s) of the model's observables via the sine-cosine tree (equivalent to the circuit expectation).

Parameters:

Name	Type	Description	Default
`params`	`Optional[ndarray]`	Model parameters. Defaults to the model's parameters.	`None`
`inputs`	`Optional[ndarray]`	Inputs to the circuit. Defaults to 1.	`None`

Returns:

Type	Description
`ndarray`	jnp.ndarray: Expectation value per observable (or their mean if `force_mean` is set).

Raises:

Type	Description
`NotImplementedError`	For execution types other than "expval" or when noise is requested.

Source code in qml_essentials/coefficients.py

def __call__(
    self,
    params: Optional[jnp.ndarray] = None,
    inputs: Optional[jnp.ndarray] = None,
    **kwargs,
) -> jnp.ndarray:
    """
    Evaluate the expectation value(s) of the model's observables via the
    sine-cosine tree (equivalent to the circuit expectation).

    Args:
        params (Optional[jnp.ndarray]): Model parameters. Defaults to the
            model's parameters.
        inputs (Optional[jnp.ndarray]): Inputs to the circuit. Defaults to 1.

    Returns:
        jnp.ndarray: Expectation value per observable (or their mean if
            ``force_mean`` is set).

    Raises:
        NotImplementedError: For execution types other than "expval" or when
            noise is requested.
    """
    params = (
        self.model._params_validation(params)
        if params is not None
        else self.model.params
    )
    inputs = (
        self.model._inputs_validation(inputs)
        if inputs is not None
        else self.model._inputs_validation(1.0)
    )

    if kwargs.get("execution_type", "expval") != "expval":
        raise NotImplementedError(
            f'Currently, only "expval" execution type is supported when '
            f"building FourierTree. Got {kwargs.get('execution_type', 'expval')}."
        )
    if kwargs.get("noise_params", None) is not None:
        raise NotImplementedError(
            "Currently, noise is not supported when building FourierTree."
        )

    # Re-derive the (canonical) parameter values for the requested inputs;
    # the tree structure (leaf arrays) is unchanged.
    operations, _ = self._build_canonical_tape(params, inputs)
    self.parameters = [
        jnp.squeeze(p) for p in PauliCircuit.get_parameters(operations)
    ]

    self._ensure_structure()
    all_columns = np.arange(self.n_params, dtype=np.int64)
    results = []
    for S, C, terms in self.leaf_arrays:
        factors = self._leaf_factors(S, C, all_columns)
        results.append(jnp.real(jnp.sum(jnp.asarray(terms) * factors)))
    results = jnp.array(results)

    if kwargs.get("force_mean", False):
        return jnp.mean(results)
    return results

`init(model)` #

Tree initialisation, based on the Pauli-Clifford representation of a model.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The Model, for which to build the tree.	required

Source code in qml_essentials/coefficients.py

def __init__(self, model: Model):
    """
    Tree initialisation, based on the Pauli-Clifford representation of a
    model.

    Args:
        model (Model): The Model, for which to build the tree.
    """
    self.model = model
    self.n_qubits = model.n_qubits

    # A single (de-batched) parameter set drives the whole tree.
    self._params = self._single_param_set(model.params)

    # Canonical Pauli-Clifford structure, recorded once at a fixed base
    # input.  The base value is irrelevant to the structure (it only sets
    # the rotation angles, not which Pauli words appear).
    base_inputs = np.ones(model.n_input_feat)
    operations, observables = self._build_canonical_tape(self._params, base_inputs)

    self.parameters = [
        jnp.squeeze(p) for p in PauliCircuit.get_parameters(operations)
    ]
    self.n_params = len(self.parameters)

    # Pauli generators of the (canonical) rotations, as symbolic words.
    self.pauli_words: List[PauliWord] = [
        PauliWord.from_operation(op, self.n_qubits) for op in operations
    ]

    # Cumulative X/Y support of the rotations[0..k] (for light-cone early
    # stopping).  cumulative_xy[k] is True on every qubit touched by an X/Y
    # generator in any rotation up to index k.
    self.cumulative_xy: List[np.ndarray] = []
    running = np.zeros(self.n_qubits, dtype=bool)
    for pw in self.pauli_words:
        running = np.logical_or(running, pw.xy_mask)
        self.cumulative_xy.append(running.copy())

    # Observable Pauli words (one tree root each).
    self.observable_words: List[PauliWord] = [
        PauliWord.from_operation(obs, self.n_qubits) for obs in observables
    ]

    # Identify the input-encoding columns, their feature, and integer
    # frequency scaling directly from the tape (no per-gate tagging).  Sets
    # ``input_indices``, ``all_input_indices``, ``input_scaling``,
    # ``var_positions`` and ``features``.
    self._detect_inputs(base_inputs)

    # The explicit leaf structure is built lazily: for deep circuits the
    # number of tree paths explodes combinatorially, while the canonical
    # form above (and the merged-state support DP) remain cheap.
    self._structure_built = False

`get_exact_support(method='tree')` #

Symbolically derive the exact frequency support (no sampling).

A frequency :math:\omega belongs to the exact spectrum iff its coefficient :math:c_\omega(\theta) = \sum_l W_{\omega l}\, \text{term}_l\, v_l(\theta) is not identically zero in the variational parameters :math:\theta.

Two methods are available:

"tree" (default, fully exact): enumerates the explicit tree leaves. Because the branch index strictly decreases along every tree path, each parameter contributes at most one sine or cosine factor per leaf (:math:S_{li}, C_{li} \in \{0, 1\}). Every variational leaf factor :math:v_l is therefore a square-free monomial over :math:\{1, \cos\theta_i, i\sin\theta_i\}, and monomials with distinct signatures are linearly independent functions (no :math:\cos^2 + \sin^2 identities can arise without squares). Hence

.. math:: c_\omega \equiv 0 \iff \sum_{l \in g} W_{\omega l}\,\text{term}_l = 0 \quad \text{for every signature group } g.

Since all involved quantities are dyadic rationals times :math:\{\pm 1, \pm i\}, the group sums are exact in float64 and the zero-test is exact. The number of leaves can however grow exponentially with circuit depth.

"dp" (scalable): merges tree nodes with identical (rotation index, observable) — at most n_params * 4^n_qubits states — and tracks, per state, the achievable per-feature sine/cosine count vectors (s_f, c_f) as a mixed-radix bitmask. Each feature's support is the union of the (exact) expansion supports of :math:\cos^{c_f} x_f\, (i \sin x_f)^{s_f}, and the model support is their Cartesian product across features. This is exact per tree path (including interior zero coefficients of the expansions), but unlike "tree" it cannot detect coefficients that cancel identically across paths with identical variational signatures (e.g. directly repeated encodings). It therefore yields a tight superset in such corner cases. Supports any number of input features, but requires unit-magnitude input scaling: per-gate :math:|\omega| \neq 1 scalings (e.g. Golomb encodings) are rejected — use "tree".

Parameters:

Name	Type	Description	Default
`method`	`str`	`"tree"` (fully exact) or `"dp"` (scalable).	`'tree'`

Returns:

Type	Description
`List[ndarray]`	List[np.ndarray]: For each observable (root), the frequency vectors
`List[ndarray]`	with not-identically-zero coefficient — shape `(n_freq,)` for a
`List[ndarray]`	single input feature, `(n_freq, n_features)` otherwise.

Source code in qml_essentials/coefficients.py

def get_exact_support(self, method: str = "tree") -> List[np.ndarray]:
    r"""Symbolically derive the exact frequency support (no sampling).

    A frequency :math:`\omega` belongs to the exact spectrum iff its
    coefficient :math:`c_\omega(\theta) = \sum_l W_{\omega l}\,
    \text{term}_l\, v_l(\theta)` is not identically zero in the
    variational parameters :math:`\theta`.

    Two methods are available:

    - ``"tree"`` (default, fully exact): enumerates the explicit tree
      leaves.  Because the branch index strictly decreases along every tree
      path, each parameter contributes **at most one** sine *or* cosine
      factor per leaf (:math:`S_{li}, C_{li} \in \{0, 1\}`).  Every
      variational leaf factor :math:`v_l` is therefore a *square-free*
      monomial over :math:`\{1, \cos\theta_i, i\sin\theta_i\}`, and
      monomials with distinct signatures are linearly independent functions
      (no :math:`\cos^2 + \sin^2` identities can arise without squares).
      Hence

      .. math::
          c_\omega \equiv 0 \iff \sum_{l \in g} W_{\omega l}\,\text{term}_l
          = 0 \quad \text{for every signature group } g.

      Since all involved quantities are dyadic rationals times
      :math:`\{\pm 1, \pm i\}`, the group sums are exact in float64 and the
      zero-test is exact.  The number of leaves can however grow
      exponentially with circuit depth.

    - ``"dp"`` (scalable): merges tree nodes with identical
      ``(rotation index, observable)`` — at most ``n_params * 4^n_qubits``
      states — and tracks, per state, the achievable per-feature sine/cosine
      count vectors ``(s_f, c_f)`` as a mixed-radix bitmask.  Each feature's
      support is the union of the (exact) expansion supports of
      :math:`\cos^{c_f} x_f\, (i \sin x_f)^{s_f}`, and the model support is
      their Cartesian product across features.  This is exact per tree path
      (including interior zero coefficients of the expansions), but unlike
      ``"tree"`` it cannot detect coefficients that cancel identically
      *across* paths with identical variational signatures (e.g. directly
      repeated encodings).  It therefore yields a tight superset in such
      corner cases.  Supports any number of input features, but requires
      unit-magnitude input scaling: per-gate :math:`|\omega| \neq 1`
      scalings (e.g. Golomb encodings) are rejected — use ``"tree"``.

    Args:
        method (str): ``"tree"`` (fully exact) or ``"dp"`` (scalable).

    Returns:
        List[np.ndarray]: For each observable (root), the frequency vectors
        with not-identically-zero coefficient — shape ``(n_freq,)`` for a
        single input feature, ``(n_freq, n_features)`` otherwise.
    """
    if method == "dp":
        return self._support_dp()
    if method != "tree":
        raise ValueError(f"Unknown method '{method}'. Use 'tree' or 'dp'.")

    self._ensure_structure()
    supports = []
    for (S, C, terms), W, freqs in zip(
        self.leaf_arrays, self.weights_per_root, self.freqs_per_root
    ):
        freqs = np.asarray(freqs)
        n_leaves = S.shape[0]
        if n_leaves == 0:
            supports.append(freqs[:0])
            continue

        # Group leaves by their variational sine/cosine signature.
        signature = np.hstack([S[:, self.var_positions], C[:, self.var_positions]])
        _, groups = np.unique(signature, axis=0, return_inverse=True)
        n_groups = int(groups.max()) + 1

        # Per-group sums of W[omega, l] * term_l, accumulated exactly.
        contrib = (W * terms[None, :]).T  # (n_leaves, n_freq)
        group_sums = np.zeros((n_groups, W.shape[0]), dtype=np.complex128)
        np.add.at(group_sums, groups, contrib)

        mask = (np.abs(group_sums) > 1e-12).any(axis=0)  # (n_freq,)
        supports.append(freqs[mask])
    return supports

`get_spectrum(force_mean=False)` #

Compute the Fourier spectrum (coefficients and frequencies) of the tree.

Parameters:

Name	Type	Description	Default
`force_mean`	`bool`	Average the coefficients over all observables (roots). Defaults to False.	`False`

Returns:

Type	Description
`Tuple[List[ndarray], List[ndarray]]`	Tuple[List[jnp.ndarray], List[jnp.ndarray]]: - List of coefficients, one entry per observable (root). - List of corresponding frequencies, one entry per root. When `force_mean` is set, both lists have a single entry.

Source code in qml_essentials/coefficients.py

def get_spectrum(
    self, force_mean: bool = False
) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
    """
    Compute the Fourier spectrum (coefficients and frequencies) of the tree.

    Args:
        force_mean (bool, optional): Average the coefficients over all
            observables (roots). Defaults to False.

    Returns:
        Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
            - List of coefficients, one entry per observable (root).
            - List of corresponding frequencies, one entry per root.
            When ``force_mean`` is set, both lists have a single entry.
    """
    self._ensure_structure()
    per_root_coeffs: List[jnp.ndarray] = []
    for (S, C, terms), W in zip(self.leaf_arrays, self.weights_per_root):
        leaf_const = jnp.asarray(terms) * self._leaf_factors(
            S, C, self.var_positions
        )
        per_root_coeffs.append(jnp.asarray(W) @ leaf_const)

    return self._combine_roots(per_root_coeffs, self.freqs_per_root, force_mean)

Fourier Coefficient Correlation#

from qml_essentials.coefficients import FCC

Source code in qml_essentials/coefficients.py

class FCC:
    @classmethod
    def get_fcc(
        cls,
        model: Model,
        n_samples: int,
        random_key: Optional[random.PRNGKey] = None,
        method: Optional[str] = "pearson",
        scale: Optional[bool] = False,
        weight: Optional[bool] = False,
        trim_redundant: Optional[bool] = True,
        **kwargs,
    ) -> float:
        """
        Shortcut method to get just the FCC.
        This includes
        1. What is done in `get_fourier_fingerprint`:
            1. Calculating the coefficients (using `n_samples`)
            2. Correlating the result from 1) using `method`
            3. Weighting the correlation matrix (if `weight` is True)
            4. Remove redundancies
        2. What is done in `calculate_fcc`:
            1. Absolute of the fingerprint
            2. Average

        Args:
            model (Model): The QFM model
            n_samples (int): Number of samples to calculate average of coefficients
            random_key (Optional[random.PRNGKey]): JAX random key for parameter
                initialization. If None, uses the model's internal random key.
            method (Optional[str], optional): Correlation method. Supported values are
                "pearson", "complex_pearson", "spearman", and "covariance".
                Defaults to "pearson".
            scale (Optional[bool], optional): Whether to scale the number of samples.
                Defaults to False.
            weight (Optional[bool], optional): Whether to weight the correlation matrix.
                Defaults to False.
            trim_redundant (Optional[bool], optional): Whether to remove redundant
                correlations. Defaults to False.
            **kwargs (Any): Additional keyword arguments for the model function.

        Returns:
            float: The FCC
        """

        # Memory-efficient fast path
        if trim_redundant and not weight:
            _, coeffs, freqs = cls._calculate_coefficients(
                model, n_samples, random_key, scale, **kwargs
            )
            pos_idx = cls._calculate_mask(freqs)
            coeffs_flat = coeffs.reshape(-1, coeffs.shape[-1])
            coeffs_sub = coeffs_flat[pos_idx]

            fp = cls._correlate(coeffs_sub.transpose(), method=method)
            abs_fp = jnp.abs(fp)
            diag = jnp.abs(jnp.diagonal(fp))

            total_sum = jnp.nansum(abs_fp)
            total_count = jnp.sum(jnp.isfinite(abs_fp))
            diag_sum = jnp.nansum(diag)
            diag_count = jnp.sum(jnp.isfinite(diag))

            lower_sum = (total_sum - diag_sum) / 2.0
            lower_count = (total_count - diag_count) / 2.0
            return lower_sum / lower_count

        fourier_fingerprint, _, _ = cls.get_fourier_fingerprint(
            model,
            n_samples,
            random_key,
            method,
            scale,
            weight,
            trim_redundant=trim_redundant,
            **kwargs,
        )

        return cls.calculate_fcc(fourier_fingerprint)

    @classmethod
    def get_fourier_fingerprint(
        cls,
        model: Model,
        n_samples: int,
        random_key: Optional[random.PRNGKey] = None,
        method: Optional[str] = "pearson",
        scale: Optional[bool] = False,
        weight: Optional[bool] = False,
        trim_redundant: Optional[bool] = True,
        nan_to_one: Optional[bool] = False,
        **kwargs: Any,
    ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
        """
        Shortcut method to get just the fourier fingerprint.
        This includes
        1. Calculating the coefficients (using `n_samples`)
        2. Correlating the result from 1) using `method`
        3. Weighting the correlation matrix (if `weight` is True)
        4. Remove redundancies (if `trim_redundant` is True)

        Args:
            model (Model): The QFM model
            n_samples (int): Number of samples to calculate average of coefficients
            random_key (Optional[random.PRNGKey]): JAX random key for parameter
                initialization. If None, uses the model's internal random key.
            method (Optional[str], optional): Correlation method. Supported values are
                "pearson", "complex_pearson", "spearman", and "covariance".
                Defaults to "pearson".
            scale (Optional[bool], optional): Whether to scale the number of samples.
                Defaults to False.
            weight (Optional[bool], optional): Whether to weight the correlation matrix.
                Defaults to False.
            trim_redundant (Optional[bool], optional): Whether to remove redundant
                correlations. Defaults to True.
            nan_to_one (Optional[bool], optional): Whether to set nan to 1.
                Defaults to False.
            **kwargs: Additional keyword arguments for the model function.

        Returns:
            Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: The fourier
            fingerprint, the corresponding frequency indices and the
            corresponding coefficients. If `trim_redundant` is True the
            frequencies are returned as a `(row_freqs, col_freqs)` tuple that
            labels the two (redundancy-trimmed) matrix axes and the
            coefficients as a matching `(row_coeffs, col_coeffs)` tuple whose
            rows align with those frequencies; otherwise the full frequency
            vector and full coefficient array are returned.
        """
        _, coeffs, freqs = cls._calculate_coefficients(
            model, n_samples, random_key, scale, **kwargs
        )

        # Memory-efficient fast path
        if trim_redundant and not weight:
            pos_idx = cls._calculate_mask(freqs)
            pos_freqs = cls._flat_frequencies(freqs)[pos_idx]

            # Flatten all frequency axes; the last axis is the sample
            # axis. `_calculate_mask` returns flat indices in C order,
            # matching this reshape.
            coeffs_flat = coeffs.reshape(-1, coeffs.shape[-1])
            coeffs_sub = coeffs_flat[pos_idx]

            fourier_fingerprint = cls._correlate(coeffs_sub.transpose(), method=method)

            if nan_to_one:
                fourier_fingerprint = jnp.where(
                    jnp.isnan(fourier_fingerprint), 1.0, fourier_fingerprint
                )

            M = fourier_fingerprint.shape[0]
            lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
            fourier_fingerprint = jnp.where(
                lower_tri_mask, fourier_fingerprint, jnp.nan
            )

            row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
            col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)
            fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]

            return (
                fourier_fingerprint,
                (pos_freqs[row_mask], pos_freqs[col_mask]),
                (coeffs_sub[row_mask], coeffs_sub[col_mask]),
            )

        fourier_fingerprint = cls._correlate(coeffs.transpose(), method=method)

        if nan_to_one:
            # set nan to 1
            fourier_fingerprint[jnp.isnan(fourier_fingerprint)] = 1.0

        # perform weighting if requested
        fourier_fingerprint = (
            cls._weighting_mean(fourier_fingerprint, coeffs)
            if weight
            else fourier_fingerprint
        )

        if trim_redundant:
            pos_idx = cls._calculate_mask(freqs)
            pos_freqs = cls._flat_frequencies(freqs)[pos_idx]
            coeffs_sub = coeffs.reshape(-1, coeffs.shape[-1])[pos_idx]

            # restrict to the positive-frequency sub-block (M x M with
            # M = number of non-negative flat-frequencies) instead of
            # building a full N x N mask. This avoids the O(N^2) float
            fourier_fingerprint = fourier_fingerprint[pos_idx][:, pos_idx]

            # keep only the strict lower triangle; the rest -> nan
            M = fourier_fingerprint.shape[0]
            lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
            fourier_fingerprint = jnp.where(
                lower_tri_mask, fourier_fingerprint, jnp.nan
            )

            row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
            col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)

            fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]

            return (
                fourier_fingerprint,
                (pos_freqs[row_mask], pos_freqs[col_mask]),
                (coeffs_sub[row_mask], coeffs_sub[col_mask]),
            )

        return fourier_fingerprint, freqs, coeffs

    @classmethod
    def calculate_fcc(
        cls,
        fourier_fingerprint: jnp.ndarray,
    ) -> float:
        """
        Method to calculate the FCC based on an existing correlation matrix.
        Calculate absolute and then the average over this matrix.
        The Fingerprint can be obtained via `get_fourier_fingerprint`

        Args:
            fourier_fingerprint (jnp.ndarray): Correlation matrix of coefficients
        Returns:
            float: The FCC
        """
        # apply the mask on the fingerprint
        return jnp.nanmean(jnp.abs(fourier_fingerprint))

    @classmethod
    def _calculate_mask(cls, freqs: jnp.ndarray) -> jnp.ndarray:
        """
        Determine the flat indices of the Fourier correlation matrix
        that lie on a non-negative-frequency row/column. Together with
        the strict-lower-triangle condition (handled by the caller),
        these indices select the entries of the correlation matrix
        that survive the redundancy filter applied in
        `get_fourier_fingerprint`:

        - rows/columns whose flat frequency component is negative are
          discarded (they are the complex-conjugate redundancies of
          their positive counterparts);
        - of the remaining positive-frequency sub-block, only the
          strict lower triangle is kept (the upper triangle, including
          the diagonal, contains either duplicates from symmetry or
          self-correlations).

        Args:
            freqs (jnp.ndarray): Array of frequencies. Either a 1-D
                vector (single input feature) or a 2-D array of shape
                ``(n_input_feat, K)`` whose rows are the per-axis
                frequency vectors.

        Returns:
            jnp.ndarray: 1-D int array of flat indices selecting the
                non-negative-frequency rows/cols of the fingerprint.
        """
        freqs_arr = jnp.asarray(freqs)

        if freqs_arr.ndim == 1:
            pos_flat = freqs_arr >= 0
        else:
            # N-D case: build the per-axis non-negativity masks and
            # combine them via broadcasting (no float `jnp.outer`!),
            # then flatten to match the row-major flattening used by
            # the upstream coefficient/correlation pipeline.
            axes_pos = [freqs_arr[i] >= 0 for i in range(freqs_arr.shape[0])]
            expanded = []
            n_axes = len(axes_pos)
            for i, p in enumerate(axes_pos):
                shape = [1] * n_axes
                shape[i] = p.shape[0]
                expanded.append(p.reshape(shape))
            nd_pos = reduce(jnp.logical_and, expanded)
            pos_flat = nd_pos.flatten()

        return jnp.where(pos_flat)[0]

    @classmethod
    def _flat_frequencies(cls, freqs: jnp.ndarray) -> jnp.ndarray:
        """
        Build the per-coefficient flat frequency labels in the same
        C-order used to flatten the coefficient/correlation pipeline, so
        they can be indexed by the flat indices from `_calculate_mask`.

        Args:
            freqs (jnp.ndarray): Either a 1-D vector (single input feature)
                or a ``(n_input_feat, K)`` stack / list of per-axis frequency
                vectors (multi-dim input).

        Returns:
            jnp.ndarray: 1-D frequency vector (single input feature) or a
                ``(N, n_input_feat)`` array of per-coefficient frequency
                tuples (multi-dim input).
        """
        fa = jnp.asarray(freqs)
        if fa.ndim == 1:
            return fa
        # Multi-dim: per-axis vectors -> flat grid of frequency tuples in the
        # same C-order used by `_calculate_mask` and the coefficient reshape.
        grids = jnp.meshgrid(*[fa[i] for i in range(fa.shape[0])], indexing="ij")
        return jnp.stack(grids, axis=-1).reshape(-1, fa.shape[0])

    @classmethod
    def _calculate_coefficients(
        cls,
        model: Model,
        n_samples: int,
        random_key: Optional[random.PRNGKey] = None,
        scale: bool = False,
        **kwargs: Any,
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """
        Calculates the Fourier coefficients of a given model
        using `n_samples`.
        Optionally, `noise_params` can be passed to perform noisy simulation.

        Args:
            model (Model): The QFM model
            n_samples (int): Number of samples to calculate average of coefficients
            random_key (Optional[random.PRNGKey]): JAX random key for parameter
                initialization. If None, uses the model's internal random key.
            scale (bool, optional): Whether to scale the number of samples.
                Defaults to False.
            **kwargs: Additional keyword arguments for the model function.

        Returns:
            Tuple[jnp.ndarray, jnp.ndarray]: Parameters and Coefficients of size NxK
        """
        if n_samples > 0:
            if scale:
                total_samples = int(
                    jnp.power(2, model.n_qubits) * n_samples * model.n_input_feat
                )
                log.info(f"Using {total_samples} samples.")
            else:
                total_samples = n_samples
            model.initialize_params(random_key, repeat=total_samples)
        else:
            total_samples = 1

        coeffs, freqs = Coefficients.get_spectrum(
            model, shift=True, trim=True, **kwargs
        )

        return model.params, coeffs, freqs

    @classmethod
    def _correlate(cls, mat: jnp.ndarray, method: str = "pearson") -> jnp.ndarray:
        """
        Correlates two arrays using `method`.
        Currently, `pearson`, `complex_pearson`, `spearman`, and `covariance`
        are supported.

        Args:
            mat (jnp.ndarray): Array of shape (N, K)
            method (str, optional): Correlation method. Defaults to "pearson".

        Raises:
            ValueError: If the method is not supported.

        Returns:
            jnp.ndarray: Correlation matrix of `a` and `b`.
        """
        assert len(mat.shape) >= 2, "Input matrix must have at least 2 dimensions"

        # Note that for the general n-D case, we have to flatten along
        # the first axis (last one is batch).
        # Note that the order here is important so we can easily filter out
        # negative coefficients later.
        # Consider the following example: [[1,2,3],[4,5,6],[7,8,9]]
        # we want to get [1, 4, 7, 2, 5, 8, 3, 6, 9]
        # such that after correlation, all positive indexed coefficients
        # will be in the bottom right quadrant
        if method == "pearson":
            result = cls._pearson(mat.reshape(mat.shape[0], -1))
            # result = cls._pearson(mat.reshape(mat.shape[-1], -1, order="F"))
        elif method == "complex_pearson":
            result = cls._complex_pearson(mat.reshape(mat.shape[0], -1))
        elif method == "spearman":
            result = cls._spearman(mat.reshape(mat.shape[0], -1))
            # result = cls._spearman(mat.reshape(mat.shape[-1], -1, order="F"))
        elif method == "covariance":
            result = cls._covariance(mat.reshape(mat.shape[0], -1))
        else:
            raise ValueError(
                f"Unknown correlation method: {method}. Must be 'pearson', \
                             'complex_pearson', 'spearman' or 'covariance'."
            )

        return result

    @classmethod
    def _covariance(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
        """
        Compute the Hermitian sample covariance between columns of `mat`,
        permitting missing values (NaN or ±Inf).

        For each pair (i, j) the covariance is computed over the rows that are
        finite in both columns, as
        sum(conj(x_i - mean_i) * (x_j - mean_j)) / (nobs - 1),
        so it computes `X.conj().T @ X`.
        Real input collapses to the ordinary real sample covariance; complex
        input yields a complex matrix whose magnitude and angle carry the
        covariance strength and relative phase.


        Args:
            mat : array_like, shape (N, K)
                Input data.
            minp : int, optional
                Minimum number of paired observations required to form a
                covariance. If the number of valid pairs for (i, j) is < minp,
                the result is NaN.

        Returns:
            cov : ndarray, shape (K, K)
                Sample covariance matrix.
        """
        mat = jnp.asarray(mat)
        real_dtype = jnp.asarray(mat.real).dtype

        mask = jnp.isfinite(mat)
        fmask = mask.astype(real_dtype)
        safe = jnp.where(mask, mat, 0.0)

        nobs = fmask.T @ fmask
        nobs_safe = jnp.where(nobs > 0, nobs, 1.0)

        sum_x = safe.T @ fmask
        sum_y = fmask.T @ safe

        masked = safe * fmask
        sum_conj_xy = jnp.conj(masked).T @ masked

        sxy = sum_conj_xy - (jnp.conj(sum_x) * sum_y) / nobs_safe

        denom = jnp.where(nobs > 1, nobs - 1, jnp.nan)
        result = sxy / denom

        result = jnp.where(nobs < minp, jnp.nan, result)

        return result

    @classmethod
    def _complex_pearson(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
        """
        Compute the complex Pearson correlation between columns of `mat`,
        permitting missing values (NaN or ±Inf).

        This uses the Hermitian normalized covariance
        sum(conj(x_i - mean_i) * (x_j - mean_j)) /
        sqrt(sum(abs(x_i - mean_i)**2) * sum(abs(x_j - mean_j)**2)).
        Consequently, if column j is exp(1j * phi) times column i, then
        abs(corr[i, j]) is 1 and angle(corr[i, j]) is phi.

        Args:
            mat : array_like, shape (N, K)
                Input data.
            minp : int, optional
                Minimum number of paired observations required to form a correlation.
                If the number of valid pairs for (i, j) is < minp, the result is NaN.

        Returns:
            corr : ndarray, shape (K, K)
                Complex Pearson correlation matrix.
        """
        mat = jnp.asarray(mat)
        real_dtype = jnp.asarray(mat.real).dtype

        mask = jnp.isfinite(mat)
        fmask = mask.astype(real_dtype)
        safe = jnp.where(mask, mat, 0.0)

        nobs = fmask.T @ fmask
        nobs_safe = jnp.where(nobs > 0, nobs, 1.0)

        sum_x = safe.T @ fmask
        sum_y = fmask.T @ safe

        masked = safe * fmask
        sum_conj_xy = jnp.conj(masked).T @ masked

        safe_abs_sq = jnp.abs(safe) ** 2
        sum_abs_x2 = safe_abs_sq.T @ fmask
        sum_abs_y2 = fmask.T @ safe_abs_sq

        ssx = sum_abs_x2 - jnp.abs(sum_x) ** 2 / nobs_safe
        ssy = sum_abs_y2 - jnp.abs(sum_y) ** 2 / nobs_safe
        sxy = sum_conj_xy - (jnp.conj(sum_x) * sum_y) / nobs_safe

        denom = jnp.sqrt(ssx * ssy)
        result = jnp.where(denom > 0, sxy / denom, jnp.nan)
        magnitude = jnp.abs(result)
        result = jnp.where(magnitude > 1.0, result / magnitude, result)

        result = jnp.where(nobs < minp, jnp.nan, result)

        return result

    @classmethod
    def _pearson(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
        """
        Compute Pearson correlation between columns of `mat`,
        permitting missing values (NaN or ±Inf).

        The Pearson correlation is the normalized covariance,
        corr[i, j] = cov[i, j] / sqrt(cov[i, i] * cov[j, j]),
        so it is obtained by normalizing `_covariance` by the per-column
        standard deviations.

        If the input is complex, real and imaginary parts are stacked along
        the sample axis so that both components contribute to the correlation
        without discarding information.

        Args:
            mat : array_like, shape (N, K)
                Input data.
            minp : int, optional
                Minimum number of paired observations required to form a correlation.
                If the number of valid pairs for (i, j) is < minp, the result is NaN.

        Returns:
            corr : ndarray, shape (K, K)
                Pearson correlation matrix.
        """
        # Preserve complex information by splitting into real / imag samples.
        # After stacking the data is real, so the Hermitian `_covariance`
        # reduces to the ordinary real sample covariance.
        if jnp.iscomplexobj(mat):
            mat = jnp.concatenate([mat.real, mat.imag], axis=0)

        cov = cls._covariance(mat, minp=minp)

        # corr[i, j] = cov[i, j] / (std_i * std_j) with std_i = sqrt(cov[i, i])
        std = jnp.sqrt(jnp.diagonal(cov))
        denom = std[:, None] * std[None, :]
        result = jnp.where(denom > 0, cov / denom, jnp.nan)

        # clip numerical drift to [-1, 1]
        result = jnp.clip(jnp.real(result), -1.0, 1.0)

        return result

    @classmethod
    def _spearman(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
        """
        Based on Pandas correlation method as implemented here:
        https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/algos.pyx

        Compute Spearman correlation between columns of `mat`,
        permitting missing values (NaN or ±Inf).

        If the input is complex, real and imaginary parts are stacked along
        the sample axis so that both components contribute to the correlation
        without discarding information.

        Args:
            mat : array_like, shape (N, K)
                Input data.
            minp : int, optional
                Minimum number of paired observations required to form a correlation.
                If the number of valid pairs for (i, j) is < minp, the result is NaN.

        Returns:
            corr : ndarray, shape (K, K)
                Spearman correlation matrix.
        """
        # Preserve complex information by splitting into real / imag samples
        if jnp.iscomplexobj(mat):
            mat = jnp.concatenate([mat.real, mat.imag], axis=0)

        mat = jnp.asarray(mat)
        N, K = mat.shape

        # trivial all-NaN answer if too few rows
        if N < minp:
            return jnp.full((K, K), jnp.nan)

        # mask of finite entries
        mask = jnp.isfinite(mat)  # shape (N, K), dtype=bool

        # precompute ranks column-wise ignoring NaNs
        ranks = np.full((N, K), np.nan)
        for j in range(K):
            valid = mask[:, j]
            if valid.any():
                ranks[valid, j] = rankdata(mat[valid, j], method="average")

        ranks = jnp.asarray(ranks)

        # Vectorised Pearson on the ranks
        # Replace NaN ranks with 0; use mask to track validity.
        rank_mask = jnp.isfinite(ranks)
        safe_ranks = jnp.where(rank_mask, ranks, 0.0)

        # Pairwise valid-observation counts  (K, K)
        fmask = rank_mask.astype(ranks.dtype)
        nobs = fmask.T @ fmask

        # Pairwise sums over mutually-valid rows
        sum_x = safe_ranks.T @ fmask  # (K, K)
        sum_y = fmask.T @ safe_ranks  # (K, K)

        # Pairwise products
        masked_ranks = safe_ranks * fmask  # same as safe_ranks
        sum_xy = masked_ranks.T @ masked_ranks  # (K, K)

        safe_sq = safe_ranks**2
        sum_x2 = safe_sq.T @ fmask  # (K, K)
        sum_y2 = fmask.T @ safe_sq  # (K, K)

        nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
        ssx = sum_x2 - sum_x**2 / nobs_safe
        ssy = sum_y2 - sum_y**2 / nobs_safe
        sxy = sum_xy - (sum_x * sum_y) / nobs_safe

        denom = jnp.sqrt(ssx * ssy)
        result = jnp.where(denom > 0, sxy / denom, jnp.nan)
        result = jnp.clip(result, -1.0, 1.0)

        # Enforce minp
        result = jnp.where(nobs < minp, jnp.nan, result)

        return result

    @classmethod
    def _weighting_linear(cls, fourier_fingerprint: jnp.ndarray) -> jnp.ndarray:
        """
        Performs weighting on the given correlation matrix.
        Here, low-frequent coefficients are weighted more heavily.

        Args:
            fourier_fingerprint (jnp.ndarray): Correlation matrix
        """
        assert (
            fourier_fingerprint.shape[0] % 2 != 0
            and fourier_fingerprint.shape[1] % 2 != 0
        ), (
            "Correlation matrix must have odd dimensions. \
            Hint: use `trim` argument when calling `get_spectrum`."
        )
        assert fourier_fingerprint.shape[0] == fourier_fingerprint.shape[1], (
            "Correlation matrix must be square."
        )

        # The weight matrix produced by the previous quadrant-mirror
        # construction has a closed form: it is a "tent" sum along the
        # two axes. Concretely, with N = fourier_fingerprint.shape[0]
        # (odd) and center = N // 2,
        #     W[i, j] = u[i] + u[j]
        # where u[k] = (center - |k - center|) / (2 * center)
        # is a triangular weighting peaking at the centre (the zero
        # frequency) and decaying linearly to 0 at the spectrum edges.
        N = fourier_fingerprint.shape[0]
        center = N // 2
        k = jnp.arange(N)
        u = (center - jnp.abs(k - center)) / (2 * center)

        return fourier_fingerprint * (u[:, None] + u[None, :])

    @classmethod
    def _weighting_mean(
        cls, fourier_fingerprint: jnp.ndarray, coeffs: jnp.ndarray
    ) -> jnp.ndarray:
        """
        Performs weighting on the given correlation matrix.
        Here, we use the product of the mean of the coefficients as weights.
        This suppresses correlations where the mean of the coefficients is near zero.

        Args:
            fourier_fingerprint (jnp.ndarray): Correlation matrix
            coeffs (jnp.ndarray): Fourier coefficients
        """
        assert fourier_fingerprint.shape[0] == fourier_fingerprint.shape[1], (
            "Correlation matrix must be square."
        )
        assert len(coeffs.shape) >= 2, (
            "Coefficient matrix must contain coefficient axes and a sample axis."
        )

        coefficient_means = jnp.abs(jnp.mean(coeffs, axis=-1))
        coefficient_means = coefficient_means.T.reshape(-1)

        assert fourier_fingerprint.shape[0] == coefficient_means.shape[0], (
            "Correlation matrix size must match the number of Fourier coefficients."
        )

        # Apply the rank-1 weight w[i] * w[j] via broadcasting instead
        # of materialising an explicit `jnp.outer` N x N intermediate.
        return (
            fourier_fingerprint
            * coefficient_means[:, None]
            * coefficient_means[None, :]
        )

`calculate_fcc(fourier_fingerprint)` `classmethod` #

Method to calculate the FCC based on an existing correlation matrix. Calculate absolute and then the average over this matrix. The Fingerprint can be obtained via get_fourier_fingerprint

Parameters:

Name	Type	Description	Default
`fourier_fingerprint`	`ndarray`	Correlation matrix of coefficients	required

Returns: float: The FCC

Source code in qml_essentials/coefficients.py

@classmethod
def calculate_fcc(
    cls,
    fourier_fingerprint: jnp.ndarray,
) -> float:
    """
    Method to calculate the FCC based on an existing correlation matrix.
    Calculate absolute and then the average over this matrix.
    The Fingerprint can be obtained via `get_fourier_fingerprint`

    Args:
        fourier_fingerprint (jnp.ndarray): Correlation matrix of coefficients
    Returns:
        float: The FCC
    """
    # apply the mask on the fingerprint
    return jnp.nanmean(jnp.abs(fourier_fingerprint))

`get_fcc(model, n_samples, random_key=None, method='pearson', scale=False, weight=False, trim_redundant=True, **kwargs)` `classmethod` #

Shortcut method to get just the FCC. This includes 1. What is done in get_fourier_fingerprint: 1. Calculating the coefficients (using n_samples) 2. Correlating the result from 1) using method 3. Weighting the correlation matrix (if weight is True) 4. Remove redundancies 2. What is done in calculate_fcc: 1. Absolute of the fingerprint 2. Average

Parameters:

Name	Type	Description	Default
`model`	`Model`	The QFM model	required
`n_samples`	`int`	Number of samples to calculate average of coefficients	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`method`	`Optional[str]`	Correlation method. Supported values are "pearson", "complex_pearson", "spearman", and "covariance". Defaults to "pearson".	`'pearson'`
`scale`	`Optional[bool]`	Whether to scale the number of samples. Defaults to False.	`False`
`weight`	`Optional[bool]`	Whether to weight the correlation matrix. Defaults to False.	`False`
`trim_redundant`	`Optional[bool]`	Whether to remove redundant correlations. Defaults to False.	`True`
`**kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Name	Type	Description
`float`	`float`	The FCC

Source code in qml_essentials/coefficients.py

@classmethod
def get_fcc(
    cls,
    model: Model,
    n_samples: int,
    random_key: Optional[random.PRNGKey] = None,
    method: Optional[str] = "pearson",
    scale: Optional[bool] = False,
    weight: Optional[bool] = False,
    trim_redundant: Optional[bool] = True,
    **kwargs,
) -> float:
    """
    Shortcut method to get just the FCC.
    This includes
    1. What is done in `get_fourier_fingerprint`:
        1. Calculating the coefficients (using `n_samples`)
        2. Correlating the result from 1) using `method`
        3. Weighting the correlation matrix (if `weight` is True)
        4. Remove redundancies
    2. What is done in `calculate_fcc`:
        1. Absolute of the fingerprint
        2. Average

    Args:
        model (Model): The QFM model
        n_samples (int): Number of samples to calculate average of coefficients
        random_key (Optional[random.PRNGKey]): JAX random key for parameter
            initialization. If None, uses the model's internal random key.
        method (Optional[str], optional): Correlation method. Supported values are
            "pearson", "complex_pearson", "spearman", and "covariance".
            Defaults to "pearson".
        scale (Optional[bool], optional): Whether to scale the number of samples.
            Defaults to False.
        weight (Optional[bool], optional): Whether to weight the correlation matrix.
            Defaults to False.
        trim_redundant (Optional[bool], optional): Whether to remove redundant
            correlations. Defaults to False.
        **kwargs (Any): Additional keyword arguments for the model function.

    Returns:
        float: The FCC
    """

    # Memory-efficient fast path
    if trim_redundant and not weight:
        _, coeffs, freqs = cls._calculate_coefficients(
            model, n_samples, random_key, scale, **kwargs
        )
        pos_idx = cls._calculate_mask(freqs)
        coeffs_flat = coeffs.reshape(-1, coeffs.shape[-1])
        coeffs_sub = coeffs_flat[pos_idx]

        fp = cls._correlate(coeffs_sub.transpose(), method=method)
        abs_fp = jnp.abs(fp)
        diag = jnp.abs(jnp.diagonal(fp))

        total_sum = jnp.nansum(abs_fp)
        total_count = jnp.sum(jnp.isfinite(abs_fp))
        diag_sum = jnp.nansum(diag)
        diag_count = jnp.sum(jnp.isfinite(diag))

        lower_sum = (total_sum - diag_sum) / 2.0
        lower_count = (total_count - diag_count) / 2.0
        return lower_sum / lower_count

    fourier_fingerprint, _, _ = cls.get_fourier_fingerprint(
        model,
        n_samples,
        random_key,
        method,
        scale,
        weight,
        trim_redundant=trim_redundant,
        **kwargs,
    )

    return cls.calculate_fcc(fourier_fingerprint)

`get_fourier_fingerprint(model, n_samples, random_key=None, method='pearson', scale=False, weight=False, trim_redundant=True, nan_to_one=False, **kwargs)` `classmethod` #

Shortcut method to get just the fourier fingerprint. This includes 1. Calculating the coefficients (using n_samples) 2. Correlating the result from 1) using method 3. Weighting the correlation matrix (if weight is True) 4. Remove redundancies (if trim_redundant is True)

Parameters:

Name	Type	Description	Default
`model`	`Model`	The QFM model	required
`n_samples`	`int`	Number of samples to calculate average of coefficients	required
`random_key`	`Optional[PRNGKey]`	JAX random key for parameter initialization. If None, uses the model's internal random key.	`None`
`method`	`Optional[str]`	Correlation method. Supported values are "pearson", "complex_pearson", "spearman", and "covariance". Defaults to "pearson".	`'pearson'`
`scale`	`Optional[bool]`	Whether to scale the number of samples. Defaults to False.	`False`
`weight`	`Optional[bool]`	Whether to weight the correlation matrix. Defaults to False.	`False`
`trim_redundant`	`Optional[bool]`	Whether to remove redundant correlations. Defaults to True.	`True`
`nan_to_one`	`Optional[bool]`	Whether to set nan to 1. Defaults to False.	`False`
`**kwargs`	`Any`	Additional keyword arguments for the model function.	`{}`

Returns:

Type	Description
`ndarray`	Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: The fourier
`ndarray`	fingerprint, the corresponding frequency indices and the
`ndarray`	corresponding coefficients. If `trim_redundant` is True the
`Tuple[ndarray, ndarray, ndarray]`	frequencies are returned as a `(row_freqs, col_freqs)` tuple that
`Tuple[ndarray, ndarray, ndarray]`	labels the two (redundancy-trimmed) matrix axes and the
`Tuple[ndarray, ndarray, ndarray]`	coefficients as a matching `(row_coeffs, col_coeffs)` tuple whose
`Tuple[ndarray, ndarray, ndarray]`	rows align with those frequencies; otherwise the full frequency
`Tuple[ndarray, ndarray, ndarray]`	vector and full coefficient array are returned.

Source code in qml_essentials/coefficients.py

@classmethod
def get_fourier_fingerprint(
    cls,
    model: Model,
    n_samples: int,
    random_key: Optional[random.PRNGKey] = None,
    method: Optional[str] = "pearson",
    scale: Optional[bool] = False,
    weight: Optional[bool] = False,
    trim_redundant: Optional[bool] = True,
    nan_to_one: Optional[bool] = False,
    **kwargs: Any,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """
    Shortcut method to get just the fourier fingerprint.
    This includes
    1. Calculating the coefficients (using `n_samples`)
    2. Correlating the result from 1) using `method`
    3. Weighting the correlation matrix (if `weight` is True)
    4. Remove redundancies (if `trim_redundant` is True)

    Args:
        model (Model): The QFM model
        n_samples (int): Number of samples to calculate average of coefficients
        random_key (Optional[random.PRNGKey]): JAX random key for parameter
            initialization. If None, uses the model's internal random key.
        method (Optional[str], optional): Correlation method. Supported values are
            "pearson", "complex_pearson", "spearman", and "covariance".
            Defaults to "pearson".
        scale (Optional[bool], optional): Whether to scale the number of samples.
            Defaults to False.
        weight (Optional[bool], optional): Whether to weight the correlation matrix.
            Defaults to False.
        trim_redundant (Optional[bool], optional): Whether to remove redundant
            correlations. Defaults to True.
        nan_to_one (Optional[bool], optional): Whether to set nan to 1.
            Defaults to False.
        **kwargs: Additional keyword arguments for the model function.

    Returns:
        Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: The fourier
        fingerprint, the corresponding frequency indices and the
        corresponding coefficients. If `trim_redundant` is True the
        frequencies are returned as a `(row_freqs, col_freqs)` tuple that
        labels the two (redundancy-trimmed) matrix axes and the
        coefficients as a matching `(row_coeffs, col_coeffs)` tuple whose
        rows align with those frequencies; otherwise the full frequency
        vector and full coefficient array are returned.
    """
    _, coeffs, freqs = cls._calculate_coefficients(
        model, n_samples, random_key, scale, **kwargs
    )

    # Memory-efficient fast path
    if trim_redundant and not weight:
        pos_idx = cls._calculate_mask(freqs)
        pos_freqs = cls._flat_frequencies(freqs)[pos_idx]

        # Flatten all frequency axes; the last axis is the sample
        # axis. `_calculate_mask` returns flat indices in C order,
        # matching this reshape.
        coeffs_flat = coeffs.reshape(-1, coeffs.shape[-1])
        coeffs_sub = coeffs_flat[pos_idx]

        fourier_fingerprint = cls._correlate(coeffs_sub.transpose(), method=method)

        if nan_to_one:
            fourier_fingerprint = jnp.where(
                jnp.isnan(fourier_fingerprint), 1.0, fourier_fingerprint
            )

        M = fourier_fingerprint.shape[0]
        lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
        fourier_fingerprint = jnp.where(
            lower_tri_mask, fourier_fingerprint, jnp.nan
        )

        row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
        col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)
        fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]

        return (
            fourier_fingerprint,
            (pos_freqs[row_mask], pos_freqs[col_mask]),
            (coeffs_sub[row_mask], coeffs_sub[col_mask]),
        )

    fourier_fingerprint = cls._correlate(coeffs.transpose(), method=method)

    if nan_to_one:
        # set nan to 1
        fourier_fingerprint[jnp.isnan(fourier_fingerprint)] = 1.0

    # perform weighting if requested
    fourier_fingerprint = (
        cls._weighting_mean(fourier_fingerprint, coeffs)
        if weight
        else fourier_fingerprint
    )

    if trim_redundant:
        pos_idx = cls._calculate_mask(freqs)
        pos_freqs = cls._flat_frequencies(freqs)[pos_idx]
        coeffs_sub = coeffs.reshape(-1, coeffs.shape[-1])[pos_idx]

        # restrict to the positive-frequency sub-block (M x M with
        # M = number of non-negative flat-frequencies) instead of
        # building a full N x N mask. This avoids the O(N^2) float
        fourier_fingerprint = fourier_fingerprint[pos_idx][:, pos_idx]

        # keep only the strict lower triangle; the rest -> nan
        M = fourier_fingerprint.shape[0]
        lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
        fourier_fingerprint = jnp.where(
            lower_tri_mask, fourier_fingerprint, jnp.nan
        )

        row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
        col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)

        fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]

        return (
            fourier_fingerprint,
            (pos_freqs[row_mask], pos_freqs[col_mask]),
            (coeffs_sub[row_mask], coeffs_sub[col_mask]),
        )

    return fourier_fingerprint, freqs, coeffs

Datasets#

from qml_essentials.coefficients import Datasets

Source code in qml_essentials/coefficients.py

class Datasets:
    @classmethod
    def generate_fourier_series(
        cls,
        random_key: random.PRNGKey,
        model: Model,
        coefficients_min: float = 0.0,
        coefficients_max: float = 1.0,
        zero_centered: bool = False,
    ) -> jnp.ndarray:
        """
        Generates the Fourier series representation of a function.
        It uses the `model.frequencies` property to retrieve the frequency
        information. This ensures that the resulting Fourier series is
        compatible with the model.

        This function is capable of generating $D$-dimensional Fourier series
        (again defined by `model.n_input_feat`).
        The highest frequency $N$ is retrieved per dimension.

        Samples of the Fourier coefficients are drawn from a uniform circle.

        Args:
            random_key (random.PRNGKey): Random number key for JAX.
            model (Model): The quantum circuit model.
            coefficients_min (float, optional): Minimum value for the coefficients.
                Defaults to 0.0.
            coefficients_max (float, optional): Maximum value for the coefficients.
                Defaults to 1.0.
            zero_centered (bool, optional): Whether to zero-center the coefficients.
                Defaults to False.

        Returns:
            jnp.ndarray: Input domain samples with shape ((N,)*D, D)
            jnp.ndarray: Fourier series values with shape ((N,)*D)
            jnp.ndarray: Fourier coefficients with shape ((N,)*D)

        """
        # TODO: the following code can be considered to
        # capturing a truly random spectrum.
        # add some constraints on the spectrum, i.e. not fully

        # Note: one key observation for understanding the following code is,
        # that instead of wrapping your head around symmetries in multi-
        # dimensional coefficient matrices, one can simply look at the flattened
        # version of such a matrix and reshape later. It just works out.
        domain_samples_per_input_dim = cls.construct_domain_samples(model)

        frequencies = cls.construct_frequencies(model)

        coefficients = cls.construct_coefficients(
            random_key, model, coefficients_min, coefficients_max, zero_centered
        )

        values = cls.calculate_values(
            domain_samples_per_input_dim, frequencies, coefficients
        )

        # return all the information we have
        return [
            domain_samples_per_input_dim.reshape(*model.degree, -1),
            values.reshape(model.degree),
            coefficients.reshape(model.degree),
        ]

    @classmethod
    def construct_domain_samples(
        cls, model: Model, mts: int = 1, mfs: int = 1
    ) -> jnp.ndarray:
        """
        Builds the input-domain sample grid for the model spectrum.

        Going from $[0, 2 \\pi \\, \\mathrm{mts}]$ with the resolution required
        for the highest frequency, permuted with the input dimensionality to get
        an n-d grid of domain samples (a "coordinate system").

        The grid follows the same convention as
        `Coefficients._fourier_transform`, so a dataset built here lands on the
        bins that `Coefficients.get_spectrum` analyses with the same `mts` and
        `mfs`. A target component at $k + j/r$ has period $2 \\pi r$, hence
        `mts` should be at least $r$ to cover a full period.

        Args:
            model (Model): The quantum circuit model.
            mts (int, optional): Domain oversampling, i.e. the number of
                periods covered. Defaults to 1.
            mfs (int, optional): Frequency oversampling, i.e. the sample
                density per period. Defaults to 1.

        Returns:
            jnp.ndarray: Domain samples with shape
                (mts $\\cdot$ mfs $\\cdot$ $\\prod$ degree, n_input_feat).
        """
        return jnp.stack(
            jnp.meshgrid(
                *[
                    jnp.arange(0, 2 * mts * jnp.pi, 2 * jnp.pi / (mfs * d))
                    for d in model.degree
                ]
            )
        ).T.reshape(-1, model.n_input_feat)

    @classmethod
    def construct_frequencies(
        cls,
        model: Model,
        random_key: Optional[random.PRNGKey] = None,
        offgrid_mode: str = "none",
        offgrid_prob: float = 0.0,
        offgrid_resolution: int = 2,
    ) -> jnp.ndarray:
        """
        Builds the frequency-index grid for the model spectrum.

        This has the same shape as the domain samples returned by
        `construct_domain_samples`.

        By default the grid is the model's own comb, so the dataset is exactly
        representable. The off-grid modes move a controllable fraction of the
        components off that comb.
        Offsets are always multiples of $1/r$ for the given resolution $r$.

        Args:
            model (Model): The quantum circuit model.
            random_key (Optional[random.PRNGKey]): Random number key for JAX.
                Required unless `offgrid_mode` is "none".
            offgrid_mode (str, optional): How to displace components off the
                model comb. "none" keeps the model comb. "index" perturbs each
                frequency independently, which spans arbitrary combs that are
                in general not exactly reachable. "generator" perturbs the
                per-gate generator frequencies and rebuilds the comb as their
                Minkowski sum, which stays exactly reachable by an encoding
                pulse configuration. Defaults to "none".
            offgrid_prob (float, optional): Probability that a single component
                ("index") or generator ("generator") is displaced. Defaults to
                0.0, which reproduces the model comb in every mode. Note that
                this is the fraction of components that end up off the comb
                only in "index" mode: a sum of displaced generators can land
                back on an integer, so "generator" mode displaces noticeably
                fewer components than asked for and saturates well below one.
            offgrid_resolution (int, optional): Denominator $r$ of the offset
                grid, i.e. offsets are drawn from $\\{\\pm j/r\\}$ with
                $j = 1 \\dots r-1$. Defaults to 2, giving half-integer offsets.

        Returns:
            jnp.ndarray: Frequency indices with shape
                ($\\prod$ degree, n_input_feat).
        """
        if offgrid_mode == "none":
            frequencies = model.frequencies
        else:
            if random_key is None:
                raise ValueError(f"offgrid_mode={offgrid_mode!r} requires a random_key")
            if offgrid_resolution < 2:
                raise ValueError(
                    f"offgrid_resolution must be at least 2, "
                    f"got {offgrid_resolution}. There is no non-integer offset "
                    "on a grid of resolution 1."
                )
            if offgrid_mode == "index":
                displace = cls._displace_indices
            elif offgrid_mode == "generator":
                displace = cls._displace_generators
            else:
                raise ValueError(
                    f"Unknown offgrid_mode: {offgrid_mode!r}. Use one of "
                    "'none', 'index', 'generator'."
                )

            frequencies = []
            for i in range(model.n_input_feat):
                random_key, sub_key = random.split(random_key)
                frequencies.append(
                    displace(model, i, sub_key, offgrid_prob, offgrid_resolution)
                )

        return jnp.stack(jnp.meshgrid(*frequencies)).T.reshape(-1, model.n_input_feat)

    @classmethod
    def _offsets(
        cls,
        random_key: random.PRNGKey,
        shape: Tuple[int, ...],
        prob: float,
        resolution: int,
    ) -> jnp.ndarray:
        """
        Draws signed offsets on the $1/r$ grid, zero where not displaced.

        Args:
            random_key (random.PRNGKey): Random number key for JAX.
            shape (Tuple[int, ...]): Shape of the offset array.
            prob (float): Probability that an entry is displaced.
            resolution (int): Denominator $r$ of the offset grid.

        Returns:
            jnp.ndarray: Offsets drawn from $\\{0\\} \\cup \\{\\pm j/r\\}$ with
                $j = 1 \\dots r-1$.
        """
        move_key, magnitude_key, sign_key = random.split(random_key, 3)
        magnitude = random.randint(magnitude_key, shape, 1, resolution) / resolution
        return (
            random.bernoulli(move_key, prob, shape)
            * random.rademacher(sign_key, shape)
            * magnitude
        )

    @classmethod
    def _displace_indices(
        cls,
        model: Model,
        feature: int,
        random_key: random.PRNGKey,
        prob: float,
        resolution: int,
    ) -> jnp.ndarray:
        """
        Displaces individual frequencies of one input feature off the comb.

        Each positive frequency is displaced independently, the result is
        re-sorted and mirrored so that the comb stays antisymmetric. This is
        what `construct_coefficients` relies on to enforce conjugate symmetry,
        and in turn what keeps the series real-valued. The comb never leaves
        the model's frequency range: an offset that would push a component past
        the highest frequency has its sign flipped rather than being clipped,
        which would put the component back on the comb.

        Args:
            model (Model): The quantum circuit model.
            feature (int): Index of the input feature.
            random_key (random.PRNGKey): Random number key for JAX.
            prob (float): Probability that a component is displaced.
            resolution (int): Denominator $r$ of the offset grid.

        Returns:
            jnp.ndarray: Displaced comb, same size as the model comb.
        """
        nominal = jnp.asarray(model.frequencies[feature])
        positive = nominal[nominal > 0]
        limit = positive[-1]

        offsets = cls._offsets(random_key, positive.shape, prob, resolution)
        # the smallest positive frequency is 1 and offsets are below 1, so only
        # the upper end of the range can be overshot
        offsets = jnp.where(positive + offsets > limit, -offsets, offsets)

        # ponytail: two components can collide (1 + 0.5 and 2 - 0.5), in which
        # case their coefficients simply add. Deduplicating would change the
        # number of components, which is the one thing the study holds fixed.
        positive = jnp.sort(positive + offsets)

        return jnp.concatenate([-jnp.flip(positive), jnp.zeros(1), positive])

    @classmethod
    def _displace_generators(
        cls,
        model: Model,
        feature: int,
        random_key: random.PRNGKey,
        prob: float,
        resolution: int,
    ) -> jnp.ndarray:
        """
        Displaces the generator frequencies of one input feature off the comb.

        Mirrors `Encoding.get_spectrum`, but scales each encoding gate's
        generator by a displaced $\\eta$ before taking the Minkowski sum, which
        is exactly what an encoding pulse scaler does to the gate it drives.
        The reachable comb then grows past the model degree in both count and
        range, so each model frequency claims the closest reachable one that is
        still inside the model's range. Staying in range matters: a component
        beyond the highest model frequency would be unreachable.
        Two model frequencies may end up claiming the same
        reachable one, as in `_displace_indices`.

        Args:
            model (Model): The quantum circuit model.
            feature (int): Index of the input feature.
            random_key (random.PRNGKey): Random number key for JAX.
            prob (float): Probability that a generator is displaced.
            resolution (int): Denominator $r$ of the offset grid.

        Returns:
            jnp.ndarray: Displaced comb, same size as the model comb.
        """
        # the offset draw and in-range flip live in _generator_etas, so the
        # scalers exposed by generator_etas cannot desync from this comb
        eta = cls._generator_etas(model, feature, random_key, prob, resolution)

        base = {"hamming": 1, "binary": 2, "ternary": 3}[model._enc._strategy]
        nominal = np.asarray(model.frequencies[feature])
        limit = nominal.max()
        mask = np.asarray(model.data_reupload[..., feature], dtype=bool)
        scale = base ** np.arange(mask.shape[1])

        # Minkowski sum over the displaced per-gate generators. Rounded before
        # deduplication, which is exact for a power-of-two resolution.
        reach = {0.0}
        for layer, qubit in zip(*np.nonzero(mask)):
            generator = scale[qubit] * eta[layer, qubit]
            reach = {
                round(a + s * generator, 9) for a in reach for s in (-1.0, 0.0, 1.0)
            }
        reachable = sorted(v for v in reach if 0 < v <= limit)

        # claim the closest reachable frequency, so the displaced comb tracks
        # the original one
        positive = jnp.sort(
            jnp.asarray(
                [
                    min(reachable, key=lambda v: abs(v - frequency))
                    for frequency in nominal[nominal > 0]
                ],
                dtype=float,
            )
        )

        return jnp.concatenate([-jnp.flip(positive), jnp.zeros(1), positive])

    @classmethod
    def _generator_etas(
        cls,
        model: Model,
        feature: int,
        random_key: random.PRNGKey,
        prob: float,
        resolution: int,
    ) -> np.ndarray:
        """
        The per-gate scalers $\\eta = 1 + \\text{offset}$ applied to the
        encoding generators of one input feature in `offgrid_mode='generator'`.

        This is the offset draw and in-range flip shared with
        `_displace_generators`; the returned array has the shape of the
        data-reupload mask ($n_\\text{layers}, n_\\text{qubits}$) and holds
        exactly the encoding pulse amplitude scalers that make the generator
        comb reachable.

        Args:
            model (Model): The quantum circuit model.
            feature (int): Index of the input feature.
            random_key (random.PRNGKey): Random number key for JAX.
            prob (float): Probability that a generator is displaced.
            resolution (int): Denominator $r$ of the offset grid.

        Returns:
            np.ndarray: Amplitude scalers $\\eta$, shape
                ($n_\\text{layers}, n_\\text{qubits}$).
        """
        base = {"hamming": 1, "binary": 2, "ternary": 3}.get(model._enc._strategy)
        if base is None:
            raise ValueError(
                f"offgrid_mode='generator' does not support the "
                f"{model._enc._strategy!r} encoding strategy, which has no "
                "per-gate pulse parametrization to displace."
            )

        nominal = np.asarray(model.frequencies[feature])
        limit = nominal.max()
        mask = np.asarray(model.data_reupload[..., feature], dtype=bool)
        scale = base ** np.arange(mask.shape[1])
        offsets = np.asarray(cls._offsets(random_key, mask.shape, prob, resolution))
        offsets = np.where(scale * (1.0 + offsets) > limit, -offsets, offsets)
        return 1.0 + offsets

    @classmethod
    def generator_etas(
        cls,
        model: Model,
        random_key: random.PRNGKey,
        offgrid_prob: float,
        offgrid_resolution: int,
    ) -> List[np.ndarray]:
        """
        The encoding pulse amplitude scalers `construct_frequencies` applies in
        `offgrid_mode='generator'`, one ($n_\\text{layers}, n_\\text{qubits}$)
        array per input feature.

        Call with the same `random_key` passed to `construct_frequencies` to
        recover the encoding pulse configuration that makes the off-grid target
        reachable, e.g. to oracle-initialize or score trained scalers against
        it. The per-feature key split mirrors `construct_frequencies`.

        Args:
            model (Model): The quantum circuit model.
            random_key (random.PRNGKey): The key passed to
                `construct_frequencies`.
            offgrid_prob (float): Probability that a generator is displaced.
            offgrid_resolution (int): Denominator $r$ of the offset grid.

        Returns:
            List[np.ndarray]: Amplitude scalers $\\eta$ per input feature.
        """
        etas = []
        for i in range(model.n_input_feat):
            random_key, sub_key = random.split(random_key)
            etas.append(
                cls._generator_etas(model, i, sub_key, offgrid_prob, offgrid_resolution)
            )
        return etas

    @classmethod
    def construct_coefficients(
        cls,
        random_key: random.PRNGKey,
        model: Model,
        coefficients_min: float = 0.0,
        coefficients_max: float = 1.0,
        zero_centered: bool = False,
    ) -> jnp.ndarray:
        """
        Samples the conjugate-symmetric Fourier coefficient vector.

        Coefficients are drawn from a uniform circle (see `uniform_circle`).
        The offset coefficient (first entry) is either zeroed or made real,
        then the spectrum is mirrored to enforce conjugate symmetry.

        Args:
            random_key (random.PRNGKey): Random number key for JAX.
            model (Model): The quantum circuit model.
            coefficients_min (float, optional): Minimum value for the
                coefficients. Defaults to 0.0.
            coefficients_max (float, optional): Maximum value for the
                coefficients. Defaults to 1.0.
            zero_centered (bool, optional): Whether to zero-center the
                coefficients. Defaults to False.

        Returns:
            jnp.ndarray: Conjugate-symmetric coefficient vector of size
                $\\prod$ degree.
        """
        coefficients = cls.uniform_circle(
            random_key,
            low=coefficients_min,
            high=coefficients_max,
            size=math.prod(model.degree) // 2 + 1,
        )

        # zero center (first coeff = 0)
        # we can assume the first coeff is the offset, because we're dealing
        # with a non-symmetric spectrum here
        if zero_centered:
            coefficients = coefficients.at[0].set(0.0)
        else:
            coefficients = coefficients.at[0].set(coefficients[0].real)

        # ensure symmetry (here, non_negative_ is removed!),
        # giving us the full coefficients vector
        return jnp.concat(
            [
                jnp.flip(coefficients[..., 1:]).conjugate(),
                coefficients,
            ],
            axis=-1,
        )

    @classmethod
    def calculate_values(
        cls,
        domain_samples: jnp.ndarray,
        frequencies: jnp.ndarray,
        coefficients: jnp.ndarray,
    ) -> jnp.ndarray:
        """
        Evaluates the real-valued Fourier series on the domain grid.

        Vectorized version of
        $f(x) = \\sum_{n=0}^{N-1} c_n e^{i \\omega_n x}$ that takes the input
        dimension into account, normalized by the number of coefficients.

        Args:
            domain_samples (jnp.ndarray): Domain samples with shape
                (n_points, n_input_feat).
            frequencies (jnp.ndarray): Frequency indices with shape
                (n_freqs, n_input_feat).
            coefficients (jnp.ndarray): Fourier coefficients with shape
                (n_freqs,).

        Returns:
            jnp.ndarray: Real-valued Fourier series samples with shape
                (n_points,).
        """
        return jnp.real(
            (jnp.exp(1j * (domain_samples @ frequencies.T)) * coefficients).sum(axis=1)
            / coefficients.size
        )

    @classmethod
    def uniform_circle(
        cls,
        random_key: random.PRNGKey,
        size: Union[jnp.ndarray, List, int],
        low=0.0,
        high=1.0,
    ):
        """
        Random number generator for complex numbers sampled inside the unit circle

        Args:
            random_key (random.PRNGKey): Random number key for JAX.
            size (Union[jnp.ndarray, int]): Number of samples. If a 2D array is passed,
                the first dimension will be the number of dimensions.
            low (float, optional): Minimum Radius. Defaults to 0.0.
            high (float, optional): Maximum Radius. Defaults to 1.0.

        Returns
            jnp.ndarray: Array of complex numbers with shape of `size`
        """

        if isinstance(size, int):
            size = jnp.array([size])

        random_key, random_key1 = random.split(random_key)
        return jnp.sqrt(
            random.uniform(random_key, size, minval=low, maxval=high)
        ) * jnp.exp(2j * jnp.pi * random.uniform(random_key1, size))

`calculate_values(domain_samples, frequencies, coefficients)` `classmethod` #

Evaluates the real-valued Fourier series on the domain grid.

Vectorized version of \(f(x) = \sum_{n=0}^{N-1} c_n e^{i \omega_n x}\) that takes the input dimension into account, normalized by the number of coefficients.

Parameters:

Name	Type	Description	Default
`domain_samples`	`ndarray`	Domain samples with shape (n_points, n_input_feat).	required
`frequencies`	`ndarray`	Frequency indices with shape (n_freqs, n_input_feat).	required
`coefficients`	`ndarray`	Fourier coefficients with shape (n_freqs,).	required

Returns:

Type	Description
`ndarray`	jnp.ndarray: Real-valued Fourier series samples with shape (n_points,).

Source code in qml_essentials/coefficients.py

@classmethod
def calculate_values(
    cls,
    domain_samples: jnp.ndarray,
    frequencies: jnp.ndarray,
    coefficients: jnp.ndarray,
) -> jnp.ndarray:
    """
    Evaluates the real-valued Fourier series on the domain grid.

    Vectorized version of
    $f(x) = \\sum_{n=0}^{N-1} c_n e^{i \\omega_n x}$ that takes the input
    dimension into account, normalized by the number of coefficients.

    Args:
        domain_samples (jnp.ndarray): Domain samples with shape
            (n_points, n_input_feat).
        frequencies (jnp.ndarray): Frequency indices with shape
            (n_freqs, n_input_feat).
        coefficients (jnp.ndarray): Fourier coefficients with shape
            (n_freqs,).

    Returns:
        jnp.ndarray: Real-valued Fourier series samples with shape
            (n_points,).
    """
    return jnp.real(
        (jnp.exp(1j * (domain_samples @ frequencies.T)) * coefficients).sum(axis=1)
        / coefficients.size
    )

`construct_coefficients(random_key, model, coefficients_min=0.0, coefficients_max=1.0, zero_centered=False)` `classmethod` #

Samples the conjugate-symmetric Fourier coefficient vector.

Coefficients are drawn from a uniform circle (see uniform_circle). The offset coefficient (first entry) is either zeroed or made real, then the spectrum is mirrored to enforce conjugate symmetry.

Parameters:

Name	Type	Description	Default
`random_key`	`PRNGKey`	Random number key for JAX.	required
`model`	`Model`	The quantum circuit model.	required
`coefficients_min`	`float`	Minimum value for the coefficients. Defaults to 0.0.	`0.0`
`coefficients_max`	`float`	Maximum value for the coefficients. Defaults to 1.0.	`1.0`
`zero_centered`	`bool`	Whether to zero-center the coefficients. Defaults to False.	`False`

Returns:

Type	Description
`ndarray`	jnp.ndarray: Conjugate-symmetric coefficient vector of size \(\prod\) degree.

Source code in qml_essentials/coefficients.py

@classmethod
def construct_coefficients(
    cls,
    random_key: random.PRNGKey,
    model: Model,
    coefficients_min: float = 0.0,
    coefficients_max: float = 1.0,
    zero_centered: bool = False,
) -> jnp.ndarray:
    """
    Samples the conjugate-symmetric Fourier coefficient vector.

    Coefficients are drawn from a uniform circle (see `uniform_circle`).
    The offset coefficient (first entry) is either zeroed or made real,
    then the spectrum is mirrored to enforce conjugate symmetry.

    Args:
        random_key (random.PRNGKey): Random number key for JAX.
        model (Model): The quantum circuit model.
        coefficients_min (float, optional): Minimum value for the
            coefficients. Defaults to 0.0.
        coefficients_max (float, optional): Maximum value for the
            coefficients. Defaults to 1.0.
        zero_centered (bool, optional): Whether to zero-center the
            coefficients. Defaults to False.

    Returns:
        jnp.ndarray: Conjugate-symmetric coefficient vector of size
            $\\prod$ degree.
    """
    coefficients = cls.uniform_circle(
        random_key,
        low=coefficients_min,
        high=coefficients_max,
        size=math.prod(model.degree) // 2 + 1,
    )

    # zero center (first coeff = 0)
    # we can assume the first coeff is the offset, because we're dealing
    # with a non-symmetric spectrum here
    if zero_centered:
        coefficients = coefficients.at[0].set(0.0)
    else:
        coefficients = coefficients.at[0].set(coefficients[0].real)

    # ensure symmetry (here, non_negative_ is removed!),
    # giving us the full coefficients vector
    return jnp.concat(
        [
            jnp.flip(coefficients[..., 1:]).conjugate(),
            coefficients,
        ],
        axis=-1,
    )

`construct_domain_samples(model, mts=1, mfs=1)` `classmethod` #

Builds the input-domain sample grid for the model spectrum.

Going from \([0, 2 \pi \, \mathrm{mts}]\) with the resolution required for the highest frequency, permuted with the input dimensionality to get an n-d grid of domain samples (a "coordinate system").

The grid follows the same convention as Coefficients._fourier_transform, so a dataset built here lands on the bins that Coefficients.get_spectrum analyses with the same mts and mfs. A target component at \(k + j/r\) has period \(2 \pi r\), hence mts should be at least \(r\) to cover a full period.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The quantum circuit model.	required
`mts`	`int`	Domain oversampling, i.e. the number of periods covered. Defaults to 1.	`1`
`mfs`	`int`	Frequency oversampling, i.e. the sample density per period. Defaults to 1.	`1`

Returns:

Type	Description
`ndarray`	jnp.ndarray: Domain samples with shape (mts \(\cdot\) mfs \(\cdot\) \(\prod\) degree, n_input_feat).

Source code in qml_essentials/coefficients.py

@classmethod
def construct_domain_samples(
    cls, model: Model, mts: int = 1, mfs: int = 1
) -> jnp.ndarray:
    """
    Builds the input-domain sample grid for the model spectrum.

    Going from $[0, 2 \\pi \\, \\mathrm{mts}]$ with the resolution required
    for the highest frequency, permuted with the input dimensionality to get
    an n-d grid of domain samples (a "coordinate system").

    The grid follows the same convention as
    `Coefficients._fourier_transform`, so a dataset built here lands on the
    bins that `Coefficients.get_spectrum` analyses with the same `mts` and
    `mfs`. A target component at $k + j/r$ has period $2 \\pi r$, hence
    `mts` should be at least $r$ to cover a full period.

    Args:
        model (Model): The quantum circuit model.
        mts (int, optional): Domain oversampling, i.e. the number of
            periods covered. Defaults to 1.
        mfs (int, optional): Frequency oversampling, i.e. the sample
            density per period. Defaults to 1.

    Returns:
        jnp.ndarray: Domain samples with shape
            (mts $\\cdot$ mfs $\\cdot$ $\\prod$ degree, n_input_feat).
    """
    return jnp.stack(
        jnp.meshgrid(
            *[
                jnp.arange(0, 2 * mts * jnp.pi, 2 * jnp.pi / (mfs * d))
                for d in model.degree
            ]
        )
    ).T.reshape(-1, model.n_input_feat)

`construct_frequencies(model, random_key=None, offgrid_mode='none', offgrid_prob=0.0, offgrid_resolution=2)` `classmethod` #

Builds the frequency-index grid for the model spectrum.

This has the same shape as the domain samples returned by construct_domain_samples.

By default the grid is the model's own comb, so the dataset is exactly representable. The off-grid modes move a controllable fraction of the components off that comb. Offsets are always multiples of \(1/r\) for the given resolution \(r\).

Parameters:

Name	Type	Description	Default
`model`	`Model`	The quantum circuit model.	required
`random_key`	`Optional[PRNGKey]`	Random number key for JAX. Required unless `offgrid_mode` is "none".	`None`
`offgrid_mode`	`str`	How to displace components off the model comb. "none" keeps the model comb. "index" perturbs each frequency independently, which spans arbitrary combs that are in general not exactly reachable. "generator" perturbs the per-gate generator frequencies and rebuilds the comb as their Minkowski sum, which stays exactly reachable by an encoding pulse configuration. Defaults to "none".	`'none'`
`offgrid_prob`	`float`	Probability that a single component ("index") or generator ("generator") is displaced. Defaults to 0.0, which reproduces the model comb in every mode. Note that this is the fraction of components that end up off the comb only in "index" mode: a sum of displaced generators can land back on an integer, so "generator" mode displaces noticeably fewer components than asked for and saturates well below one.	`0.0`
`offgrid_resolution`	`int`	Denominator \(r\) of the offset grid, i.e. offsets are drawn from \(\{\pm j/r\}\) with \(j = 1 \dots r-1\). Defaults to 2, giving half-integer offsets.	`2`

Returns:

Type	Description
`ndarray`	jnp.ndarray: Frequency indices with shape (\(\prod\) degree, n_input_feat).

Source code in qml_essentials/coefficients.py

@classmethod
def construct_frequencies(
    cls,
    model: Model,
    random_key: Optional[random.PRNGKey] = None,
    offgrid_mode: str = "none",
    offgrid_prob: float = 0.0,
    offgrid_resolution: int = 2,
) -> jnp.ndarray:
    """
    Builds the frequency-index grid for the model spectrum.

    This has the same shape as the domain samples returned by
    `construct_domain_samples`.

    By default the grid is the model's own comb, so the dataset is exactly
    representable. The off-grid modes move a controllable fraction of the
    components off that comb.
    Offsets are always multiples of $1/r$ for the given resolution $r$.

    Args:
        model (Model): The quantum circuit model.
        random_key (Optional[random.PRNGKey]): Random number key for JAX.
            Required unless `offgrid_mode` is "none".
        offgrid_mode (str, optional): How to displace components off the
            model comb. "none" keeps the model comb. "index" perturbs each
            frequency independently, which spans arbitrary combs that are
            in general not exactly reachable. "generator" perturbs the
            per-gate generator frequencies and rebuilds the comb as their
            Minkowski sum, which stays exactly reachable by an encoding
            pulse configuration. Defaults to "none".
        offgrid_prob (float, optional): Probability that a single component
            ("index") or generator ("generator") is displaced. Defaults to
            0.0, which reproduces the model comb in every mode. Note that
            this is the fraction of components that end up off the comb
            only in "index" mode: a sum of displaced generators can land
            back on an integer, so "generator" mode displaces noticeably
            fewer components than asked for and saturates well below one.
        offgrid_resolution (int, optional): Denominator $r$ of the offset
            grid, i.e. offsets are drawn from $\\{\\pm j/r\\}$ with
            $j = 1 \\dots r-1$. Defaults to 2, giving half-integer offsets.

    Returns:
        jnp.ndarray: Frequency indices with shape
            ($\\prod$ degree, n_input_feat).
    """
    if offgrid_mode == "none":
        frequencies = model.frequencies
    else:
        if random_key is None:
            raise ValueError(f"offgrid_mode={offgrid_mode!r} requires a random_key")
        if offgrid_resolution < 2:
            raise ValueError(
                f"offgrid_resolution must be at least 2, "
                f"got {offgrid_resolution}. There is no non-integer offset "
                "on a grid of resolution 1."
            )
        if offgrid_mode == "index":
            displace = cls._displace_indices
        elif offgrid_mode == "generator":
            displace = cls._displace_generators
        else:
            raise ValueError(
                f"Unknown offgrid_mode: {offgrid_mode!r}. Use one of "
                "'none', 'index', 'generator'."
            )

        frequencies = []
        for i in range(model.n_input_feat):
            random_key, sub_key = random.split(random_key)
            frequencies.append(
                displace(model, i, sub_key, offgrid_prob, offgrid_resolution)
            )

    return jnp.stack(jnp.meshgrid(*frequencies)).T.reshape(-1, model.n_input_feat)

`generate_fourier_series(random_key, model, coefficients_min=0.0, coefficients_max=1.0, zero_centered=False)` `classmethod` #

Generates the Fourier series representation of a function. It uses the model.frequencies property to retrieve the frequency information. This ensures that the resulting Fourier series is compatible with the model.

This function is capable of generating \(D\)-dimensional Fourier series (again defined by model.n_input_feat). The highest frequency \(N\) is retrieved per dimension.

Samples of the Fourier coefficients are drawn from a uniform circle.

Parameters:

Name	Type	Description	Default
`random_key`	`PRNGKey`	Random number key for JAX.	required
`model`	`Model`	The quantum circuit model.	required
`coefficients_min`	`float`	Minimum value for the coefficients. Defaults to 0.0.	`0.0`
`coefficients_max`	`float`	Maximum value for the coefficients. Defaults to 1.0.	`1.0`
`zero_centered`	`bool`	Whether to zero-center the coefficients. Defaults to False.	`False`

Returns:

Type	Description
`ndarray`	jnp.ndarray: Input domain samples with shape ((N,)*D, D)
`ndarray`	jnp.ndarray: Fourier series values with shape ((N,)*D)
`ndarray`	jnp.ndarray: Fourier coefficients with shape ((N,)*D)

Source code in qml_essentials/coefficients.py

@classmethod
def generate_fourier_series(
    cls,
    random_key: random.PRNGKey,
    model: Model,
    coefficients_min: float = 0.0,
    coefficients_max: float = 1.0,
    zero_centered: bool = False,
) -> jnp.ndarray:
    """
    Generates the Fourier series representation of a function.
    It uses the `model.frequencies` property to retrieve the frequency
    information. This ensures that the resulting Fourier series is
    compatible with the model.

    This function is capable of generating $D$-dimensional Fourier series
    (again defined by `model.n_input_feat`).
    The highest frequency $N$ is retrieved per dimension.

    Samples of the Fourier coefficients are drawn from a uniform circle.

    Args:
        random_key (random.PRNGKey): Random number key for JAX.
        model (Model): The quantum circuit model.
        coefficients_min (float, optional): Minimum value for the coefficients.
            Defaults to 0.0.
        coefficients_max (float, optional): Maximum value for the coefficients.
            Defaults to 1.0.
        zero_centered (bool, optional): Whether to zero-center the coefficients.
            Defaults to False.

    Returns:
        jnp.ndarray: Input domain samples with shape ((N,)*D, D)
        jnp.ndarray: Fourier series values with shape ((N,)*D)
        jnp.ndarray: Fourier coefficients with shape ((N,)*D)

    """
    # TODO: the following code can be considered to
    # capturing a truly random spectrum.
    # add some constraints on the spectrum, i.e. not fully

    # Note: one key observation for understanding the following code is,
    # that instead of wrapping your head around symmetries in multi-
    # dimensional coefficient matrices, one can simply look at the flattened
    # version of such a matrix and reshape later. It just works out.
    domain_samples_per_input_dim = cls.construct_domain_samples(model)

    frequencies = cls.construct_frequencies(model)

    coefficients = cls.construct_coefficients(
        random_key, model, coefficients_min, coefficients_max, zero_centered
    )

    values = cls.calculate_values(
        domain_samples_per_input_dim, frequencies, coefficients
    )

    # return all the information we have
    return [
        domain_samples_per_input_dim.reshape(*model.degree, -1),
        values.reshape(model.degree),
        coefficients.reshape(model.degree),
    ]

`generator_etas(model, random_key, offgrid_prob, offgrid_resolution)` `classmethod` #

The encoding pulse amplitude scalers construct_frequencies applies in offgrid_mode='generator', one (\(n_\text{layers}, n_\text{qubits}\)) array per input feature.

Call with the same random_key passed to construct_frequencies to recover the encoding pulse configuration that makes the off-grid target reachable, e.g. to oracle-initialize or score trained scalers against it. The per-feature key split mirrors construct_frequencies.

Parameters:

Name	Type	Description	Default
`model`	`Model`	The quantum circuit model.	required
`random_key`	`PRNGKey`	The key passed to `construct_frequencies`.	required
`offgrid_prob`	`float`	Probability that a generator is displaced.	required
`offgrid_resolution`	`int`	Denominator \(r\) of the offset grid.	required

Returns:

Type	Description
`List[ndarray]`	List[np.ndarray]: Amplitude scalers \(\eta\) per input feature.

Source code in qml_essentials/coefficients.py

@classmethod
def generator_etas(
    cls,
    model: Model,
    random_key: random.PRNGKey,
    offgrid_prob: float,
    offgrid_resolution: int,
) -> List[np.ndarray]:
    """
    The encoding pulse amplitude scalers `construct_frequencies` applies in
    `offgrid_mode='generator'`, one ($n_\\text{layers}, n_\\text{qubits}$)
    array per input feature.

    Call with the same `random_key` passed to `construct_frequencies` to
    recover the encoding pulse configuration that makes the off-grid target
    reachable, e.g. to oracle-initialize or score trained scalers against
    it. The per-feature key split mirrors `construct_frequencies`.

    Args:
        model (Model): The quantum circuit model.
        random_key (random.PRNGKey): The key passed to
            `construct_frequencies`.
        offgrid_prob (float): Probability that a generator is displaced.
        offgrid_resolution (int): Denominator $r$ of the offset grid.

    Returns:
        List[np.ndarray]: Amplitude scalers $\\eta$ per input feature.
    """
    etas = []
    for i in range(model.n_input_feat):
        random_key, sub_key = random.split(random_key)
        etas.append(
            cls._generator_etas(model, i, sub_key, offgrid_prob, offgrid_resolution)
        )
    return etas

`uniform_circle(random_key, size, low=0.0, high=1.0)` `classmethod` #

Random number generator for complex numbers sampled inside the unit circle

Parameters:

Name	Type	Description	Default
`random_key`	`PRNGKey`	Random number key for JAX.	required
`size`	`Union[ndarray, int]`	Number of samples. If a 2D array is passed, the first dimension will be the number of dimensions.	required
`low`	`float`	Minimum Radius. Defaults to 0.0.	`0.0`
`high`	`float`	Maximum Radius. Defaults to 1.0.	`1.0`

Returns jnp.ndarray: Array of complex numbers with shape of size

Source code in qml_essentials/coefficients.py

@classmethod
def uniform_circle(
    cls,
    random_key: random.PRNGKey,
    size: Union[jnp.ndarray, List, int],
    low=0.0,
    high=1.0,
):
    """
    Random number generator for complex numbers sampled inside the unit circle

    Args:
        random_key (random.PRNGKey): Random number key for JAX.
        size (Union[jnp.ndarray, int]): Number of samples. If a 2D array is passed,
            the first dimension will be the number of dimensions.
        low (float, optional): Minimum Radius. Defaults to 0.0.
        high (float, optional): Maximum Radius. Defaults to 1.0.

    Returns
        jnp.ndarray: Array of complex numbers with shape of `size`
    """

    if isinstance(size, int):
        size = jnp.array([size])

    random_key, random_key1 = random.split(random_key)
    return jnp.sqrt(
        random.uniform(random_key, size, minval=low, maxval=high)
    ) * jnp.exp(2j * jnp.pi * random.uniform(random_key1, size))

Topologies#

from qml_essentials.topologies import Topology

Generates [control, target] wire-pair lists for two-qubit gates.

All public methods are static and share a small set of private helpers so that related topologies (e.g. linear / circular, brick_layer / brick_layer_wrap) re-use the same core logic.

Raises#

ValueError If n_qubits < 2 is passed to any topology method.

Source code in qml_essentials/topologies.py

class Topology:
    """
    Generates [control, target] wire-pair lists for two-qubit gates.

    All public methods are static and share a small set of private
    helpers so that related topologies (e.g. ``linear`` / ``circular``,
    ``brick_layer`` / ``brick_layer_wrap``) re-use the same core logic.

    Raises
    ------
    ValueError
        If ``n_qubits < 2`` is passed to any topology method.
    """

    @classmethod
    def stairs(
        cls,
        n_qubits: int,
        offset: Union[int, Callable] = 0,
        wrap=False,
        reverse: bool = True,
        mirror: bool = True,
        span: Union[int, Callable] = 1,
        stride: int = 1,
        modulo: bool = True,
    ) -> List[List[int]]:
        """
        Unified generator for nearest-neighbour and spand pair topologies.
        Produces ``[control, target]`` pairs of qubits.

        The default values, produce an "upstairs" entangling sequence
        without wrapping around the last gate.

        Parameters
        ----------
        n_qubits : int
            Number of qubits.
        offset : Union[int, Callable]
            Offset for starting the entangling sequence.
            Can either be a integer or a callable that takes n_qubits as input.
        wrap : bool
            Wraps around the entangling gates.
        reverse : bool
            Reverses both the iteration direction (upstairs/ downstairs)
        mirror: bool
            Flip target/ control qubit
        span : int
            Offset between control and target qubit. Defaults to 1
        stride : int
            Step size for entangling gates. Defaults to 1, meaning a stair
            pattern will be generated.
        modulo : bool
            If a gate should be placed when the iterator decreases below 0
            or exceeds n_qubits. Defaults to True

        Returns
        -------
        List[List[int]]
        """
        ctrls = []
        targets = []

        n_gates = n_qubits if wrap else n_qubits - 1
        _offset = offset(n_qubits) if callable(offset) else offset
        _span = span(n_qubits) if callable(span) else span

        for q in range(0, n_gates, stride):
            _target = q + _offset + _span
            if _target >= n_qubits and not modulo:
                continue
            _control = q + _offset
            if _control < 0 and not modulo:
                continue

            _target = _target % n_qubits
            _control = _control % n_qubits

            if _target == _control:
                log.warning("Skipping gate where control == target")
                continue

            targets += [_target]
            ctrls += [_control]

        if reverse:
            ctrls = reversed(ctrls)
            targets = reversed(targets)

        if mirror:
            ctrls, targets = targets, ctrls

        pairs = list(zip(ctrls, targets, strict=True))

        return pairs

    @classmethod
    def bricks(cls, n_qubits: int, **kwargs) -> List[List[int]]:
        kwargs.setdefault("stride", 2)
        kwargs.setdefault("modulo", False)
        return cls.stairs(n_qubits=n_qubits, **kwargs)

    @classmethod
    def graph(
        cls, n_qubits: int, *, edges: Sequence[Sequence[int]]
    ) -> List[Tuple[int, int]]:
        """
        Explicit edge list as a topology.

        The given order and orientation are preserved, so the resulting
        circuit is deterministic and directed gates act on the wires as
        written. Both orientations of the same qubit pair are therefore
        allowed; only a repeated ``(control, target)`` pair is rejected.

        Parameters
        ----------
        n_qubits : int
            Number of qubits.
        edges : Sequence[Sequence[int]]
            ``(control, target)`` qubit pairs.

        Returns
        -------
        List[Tuple[int, int]]

        Raises
        ------
        ValueError
            If an edge leaves the qubit range, is a self-loop or repeats.
        """
        seen = set()
        pairs = []
        for q, r in edges:
            if not (0 <= q < n_qubits and 0 <= r < n_qubits) or q == r:
                raise ValueError(f"edge ({q}, {r}) invalid on {n_qubits} qubits")
            if (q, r) in seen:
                raise ValueError(f"duplicate edge ({q}, {r})")
            seen.add((q, r))
            pairs.append((q, r))

        return pairs

    @classmethod
    def all_pairs(cls, n_qubits: int) -> List[List[int]]:
        """Every unordered pair ``[j, k]`` with ``j < k``."""
        return [[j, k] for j, k in combinations(range(n_qubits), 2)]

    @classmethod
    def all_to_all(cls, n_qubits: int) -> List[List[int]]:
        """Every ordered pair ``(i, j)`` with ``i ≠ j``."""
        pairs: List[List[int]] = []
        for ql in range(n_qubits):
            for q in range(n_qubits):
                if q != ql:
                    pairs.append(
                        [
                            n_qubits - ql - 1,
                            (n_qubits - q - 1) % n_qubits,
                        ]
                    )
        return pairs

`all_pairs(n_qubits)` `classmethod` #

Every unordered pair [j, k] with j < k.

Source code in qml_essentials/topologies.py

@classmethod
def all_pairs(cls, n_qubits: int) -> List[List[int]]:
    """Every unordered pair ``[j, k]`` with ``j < k``."""
    return [[j, k] for j, k in combinations(range(n_qubits), 2)]

`all_to_all(n_qubits)` `classmethod` #

Every ordered pair (i, j) with i ≠ j.

Source code in qml_essentials/topologies.py

@classmethod
def all_to_all(cls, n_qubits: int) -> List[List[int]]:
    """Every ordered pair ``(i, j)`` with ``i ≠ j``."""
    pairs: List[List[int]] = []
    for ql in range(n_qubits):
        for q in range(n_qubits):
            if q != ql:
                pairs.append(
                    [
                        n_qubits - ql - 1,
                        (n_qubits - q - 1) % n_qubits,
                    ]
                )
    return pairs

`graph(n_qubits, *, edges)` `classmethod` #

Explicit edge list as a topology.

The given order and orientation are preserved, so the resulting circuit is deterministic and directed gates act on the wires as written. Both orientations of the same qubit pair are therefore allowed; only a repeated (control, target) pair is rejected.

Parameters#

n_qubits : int Number of qubits. edges : Sequence[Sequence[int]] (control, target) qubit pairs.

Returns#

List[Tuple[int, int]]

Raises#

ValueError If an edge leaves the qubit range, is a self-loop or repeats.

Source code in qml_essentials/topologies.py

@classmethod
def graph(
    cls, n_qubits: int, *, edges: Sequence[Sequence[int]]
) -> List[Tuple[int, int]]:
    """
    Explicit edge list as a topology.

    The given order and orientation are preserved, so the resulting
    circuit is deterministic and directed gates act on the wires as
    written. Both orientations of the same qubit pair are therefore
    allowed; only a repeated ``(control, target)`` pair is rejected.

    Parameters
    ----------
    n_qubits : int
        Number of qubits.
    edges : Sequence[Sequence[int]]
        ``(control, target)`` qubit pairs.

    Returns
    -------
    List[Tuple[int, int]]

    Raises
    ------
    ValueError
        If an edge leaves the qubit range, is a self-loop or repeats.
    """
    seen = set()
    pairs = []
    for q, r in edges:
        if not (0 <= q < n_qubits and 0 <= r < n_qubits) or q == r:
            raise ValueError(f"edge ({q}, {r}) invalid on {n_qubits} qubits")
        if (q, r) in seen:
            raise ValueError(f"duplicate edge ({q}, {r})")
        seen.add((q, r))
        pairs.append((q, r))

    return pairs

`stairs(n_qubits, offset=0, wrap=False, reverse=True, mirror=True, span=1, stride=1, modulo=True)` `classmethod` #

Unified generator for nearest-neighbour and spand pair topologies. Produces [control, target] pairs of qubits.

The default values, produce an "upstairs" entangling sequence without wrapping around the last gate.

Parameters#

n_qubits : int Number of qubits. offset : Union[int, Callable] Offset for starting the entangling sequence. Can either be a integer or a callable that takes n_qubits as input. wrap : bool Wraps around the entangling gates. reverse : bool Reverses both the iteration direction (upstairs/ downstairs) mirror: bool Flip target/ control qubit span : int Offset between control and target qubit. Defaults to 1 stride : int Step size for entangling gates. Defaults to 1, meaning a stair pattern will be generated. modulo : bool If a gate should be placed when the iterator decreases below 0 or exceeds n_qubits. Defaults to True

Returns#

List[List[int]]

Source code in qml_essentials/topologies.py

@classmethod
def stairs(
    cls,
    n_qubits: int,
    offset: Union[int, Callable] = 0,
    wrap=False,
    reverse: bool = True,
    mirror: bool = True,
    span: Union[int, Callable] = 1,
    stride: int = 1,
    modulo: bool = True,
) -> List[List[int]]:
    """
    Unified generator for nearest-neighbour and spand pair topologies.
    Produces ``[control, target]`` pairs of qubits.

    The default values, produce an "upstairs" entangling sequence
    without wrapping around the last gate.

    Parameters
    ----------
    n_qubits : int
        Number of qubits.
    offset : Union[int, Callable]
        Offset for starting the entangling sequence.
        Can either be a integer or a callable that takes n_qubits as input.
    wrap : bool
        Wraps around the entangling gates.
    reverse : bool
        Reverses both the iteration direction (upstairs/ downstairs)
    mirror: bool
        Flip target/ control qubit
    span : int
        Offset between control and target qubit. Defaults to 1
    stride : int
        Step size for entangling gates. Defaults to 1, meaning a stair
        pattern will be generated.
    modulo : bool
        If a gate should be placed when the iterator decreases below 0
        or exceeds n_qubits. Defaults to True

    Returns
    -------
    List[List[int]]
    """
    ctrls = []
    targets = []

    n_gates = n_qubits if wrap else n_qubits - 1
    _offset = offset(n_qubits) if callable(offset) else offset
    _span = span(n_qubits) if callable(span) else span

    for q in range(0, n_gates, stride):
        _target = q + _offset + _span
        if _target >= n_qubits and not modulo:
            continue
        _control = q + _offset
        if _control < 0 and not modulo:
            continue

        _target = _target % n_qubits
        _control = _control % n_qubits

        if _target == _control:
            log.warning("Skipping gate where control == target")
            continue

        targets += [_target]
        ctrls += [_control]

    if reverse:
        ctrls = reversed(ctrls)
        targets = reversed(targets)

    if mirror:
        ctrls, targets = targets, ctrls

    pairs = list(zip(ctrls, targets, strict=True))

    return pairs

Pauli Circuit#

from qml_essentials.pauli import PauliCircuit

Wrapper for Pauli-Clifford Circuits described by Nemkov et al. (https://doi.org/10.1103/PhysRevA.108.032406). The code is inspired by the corresponding implementation: https://github.com/idnm/FourierVQA.

A Pauli Circuit only consists of parameterised Pauli-rotations and Clifford gates, which is the default for the most common VQCs.

Source code in qml_essentials/pauli.py

class PauliCircuit:
    """
    Wrapper for Pauli-Clifford Circuits described by Nemkov et al.
    (https://doi.org/10.1103/PhysRevA.108.032406). The code is inspired
    by the corresponding implementation: https://github.com/idnm/FourierVQA.

    A Pauli Circuit only consists of parameterised Pauli-rotations and Clifford
    gates, which is the default for the most common VQCs.
    """

    PAULI_ROTATION_GATES = (
        RX,
        RY,
        RZ,
        PauliRot,
    )

    SKIPPABLE_OPERATIONS = (Barrier,)

    @staticmethod
    def from_parameterised_circuit(
        tape: List[Operation],
        observables: Optional[List[Operation]] = None,
        n_qubits: Optional[int] = None,
    ) -> Tuple[List[Operation], List[Operation]]:
        """
        Transforms a list of operations into a Pauli-Clifford circuit.

        Args:
            tape: List of operations recorded from the circuit.
            observables: List of observable operations.  If ``None``, defaults
                to an empty list.
            n_qubits: Total number of qubits.  Inferred from the maximum wire
                index if ``None``.

        Returns:
            Tuple[List[Operation], List[Operation]]:
                The Pauli rotations of the canonical circuit and the
                (Clifford-evolved) observables.
        """
        if observables is None:
            observables = []

        operations = PauliCircuit.get_clifford_pauli_gates(tape)

        if n_qubits is None:
            n_qubits = PauliCircuit._infer_n_qubits(operations, observables)

        pauli_gates, final_cliffords = PauliCircuit.commute_all_cliffords_to_the_end(
            operations, n_qubits
        )

        observables = PauliCircuit.cliffords_in_observable(
            final_cliffords, observables, n_qubits
        )

        return pauli_gates, observables

    @staticmethod
    def get_parameters(operations: List[Operation]) -> list:
        """Flatten the parameter values of a tape (list of operations)."""
        return [p for op in operations for p in op.parameters]

    @staticmethod
    def _infer_n_qubits(
        operations: List[Operation], observables: List[Operation]
    ) -> int:
        """Infer the register size from the maximum wire index used."""
        max_wire = -1
        for op in list(operations) + list(observables):
            if op.wires:
                max_wire = max(max_wire, max(op.wires))
        return max_wire + 1

    @staticmethod
    def commute_all_cliffords_to_the_end(
        operations: List[Operation],
        n_qubits: int,
    ) -> Tuple[List[Operation], List[Operation]]:
        """
        This function moves all clifford gates to the end of the circuit,
        accounting for commutation rules.

        Args:
            operations (List[Operation]): The operations in the tape of the
                circuit
            n_qubits (int): Total number of qubits.

        Returns:
            Tuple[List[Operation], List[Operation]]:
                - List of the resulting Pauli-rotations
                - List of the resulting Clifford gates
        """
        first_clifford = -1
        for i in range(len(operations) - 2, -1, -1):
            j = i
            while (
                j + 1 < len(operations)  # Clifford has not alredy reached the end
                and PauliCircuit._is_clifford(operations[j])
                and PauliCircuit._is_pauli_rotation(operations[j + 1])
            ):
                pauli, clifford = PauliCircuit._evolve_clifford_rotation(
                    operations[j], operations[j + 1], n_qubits
                )
                operations[j] = pauli
                operations[j + 1] = clifford
                j += 1
                first_clifford = j

        # No Clifford gates are in the circuit
        if not PauliCircuit._is_clifford(operations[-1]):
            return operations, []

        pauli_rotations = operations[:first_clifford]
        clifford_gates = operations[first_clifford:]

        return pauli_rotations, clifford_gates

    @staticmethod
    def get_clifford_pauli_gates(tape: List[Operation]) -> List[Operation]:
        """
        This function decomposes all gates in the circuit to clifford and
        pauli-rotation gates.

        Args:
            tape: List of operations recorded from the circuit.

        Returns:
            List[Operation]: A list of operations consisting only of clifford
                and Pauli-rotation gates.
        """
        operations = []
        for operation in tape:
            if PauliCircuit._is_clifford(operation) or PauliCircuit._is_pauli_rotation(
                operation
            ):
                operations.append(operation)
            elif PauliCircuit._is_skippable(operation):
                continue
            else:
                # Composite gates (Rot, CRX/CRY/CRZ, ...) expose their own
                # Clifford + Pauli-rotation decomposition.
                try:
                    operations.extend(operation.decompose())
                except NotImplementedError:
                    raise NotImplementedError(
                        f"Gate {operation.name} cannot be decomposed into "
                        "Pauli rotations and Clifford gates. Consider using a "
                        "circuit ansatz that only uses RX, RY, RZ, PauliRot, "
                        "Rot, and standard Clifford gates."
                    )

        return operations

    @staticmethod
    def _is_skippable(operation: Operation) -> bool:
        """Whether an operation can be ignored (currently only barriers)."""
        return isinstance(operation, PauliCircuit.SKIPPABLE_OPERATIONS)

    @staticmethod
    def _is_clifford(operation: Operation) -> bool:
        """Whether an operation is a Clifford gate (reads ``Operation.is_clifford``).

        Clifford gates are commuted to the end via symbolic conjugation
        (:meth:`PauliWord.conjugate_by_clifford`); see ``Operation.is_clifford``.
        """
        return getattr(operation, "is_clifford", False)

    @staticmethod
    def _is_pauli_rotation(operation: Operation) -> bool:
        """Whether an operation is a Pauli rotation gate."""
        return isinstance(operation, PauliCircuit.PAULI_ROTATION_GATES)

    @staticmethod
    def _evolve_clifford_rotation(
        clifford: Operation, pauli: Operation, n_qubits: int
    ) -> Tuple[Operation, Operation]:
        """
        Compute the resulting operations when switching a Clifford gate and a
        Pauli rotation in the circuit, i.e. move the Clifford past the rotation:

        ``... C R_P(phi) ...  ->  ... R_{C P C^dagger}(phi) C ...``

        The evolved Pauli rotation is obtained by **symbolic** Clifford
        conjugation of the rotation generator (no matrices).

        Args:
            clifford (Operation): Clifford gate to move.
            pauli (Operation): Pauli rotation gate to move the clifford past.
            n_qubits (int): Total number of qubits.

        Returns:
            Tuple[Operation, Operation]:
                - Evolved Pauli rotation operator
                - The (unchanged) Clifford operator
        """
        if not any(p_c in clifford.wires for p_c in pauli.wires):
            return pauli, clifford

        param = pauli.parameters[0]

        gen_word = PauliWord.from_operation(pauli, n_qubits)
        evolved = gen_word.conjugate_by_clifford(clifford, adjoint_left=False)
        bare, phase = evolved.to_pauli_string_and_phase()

        # Clifford conjugation of a (Hermitian) Pauli generator yields +-1.
        param_factor = float(np.real(phase))

        pauli_str, qubits = PauliCircuit._remove_identities_from_paulistr(
            bare, list(range(n_qubits))
        )
        new_pauli = PauliRot(param * param_factor, pauli_str, qubits)

        return new_pauli, clifford

    @staticmethod
    def _remove_identities_from_paulistr(
        pauli_str: str, qubits: List[int]
    ) -> Tuple[str, List[int]]:
        """
        Removes identities from Pauli string and its corresponding qubits.

        Args:
            pauli_str (str): Pauli string
            qubits (List[int]): Corresponding qubit indices.

        Returns:
            Tuple[str, List[int]]:
                - Pauli string without identities
                - Qubits indices without the identities
        """

        reduced_qubits = []
        reduced_pauli_str = ""
        for i, p in enumerate(pauli_str):
            if p != "I":
                reduced_pauli_str += p
                reduced_qubits.append(qubits[i])

        return reduced_pauli_str, reduced_qubits

    @staticmethod
    def cliffords_in_observable(
        operations: List[Operation],
        original_obs: List[Operation],
        n_qubits: int,
    ) -> List[Operation]:
        """
        Integrates Clifford gates into the observables of the original ansatz,
        by symbolically conjugating each observable through the final Clifford
        sequence (``O -> C^dagger O C`` for each Clifford, applied in reverse).

        Args:
            operations (List[Operation]): Clifford gates
            original_obs (List[Operation]): Original observables from the
                circuit
            n_qubits (int): Total number of qubits.

        Returns:
            List[Operation]: Observables with Clifford operations absorbed.
                Each carries a cached symbolic ``_pauli_word`` for the
                Fourier-tree algorithm and a matrix for simulation.
        """
        observables = []
        for ob in original_obs:
            word = PauliWord.from_operation(ob, n_qubits)
            for clifford in operations[::-1]:
                word = word.conjugate_by_clifford(clifford, adjoint_left=True)
            observables.append(PauliCircuit._pauli_operation_from_word(word))
        return observables

    @staticmethod
    def _pauli_operation_from_word(word: PauliWord) -> Operation:
        """Build an observable :class:`Operation` from a symbolic Pauli word.

        The returned operation carries both a dense ``matrix`` (for the
        statevector simulator) and a cached ``_pauli_word`` / ``_pauli_label``
        (for symbolic consumers such as the Fourier tree).
        """
        bare, phase = word.to_pauli_string_and_phase()
        reduced_str, reduced_wires = PauliCircuit._remove_identities_from_paulistr(
            bare, list(range(word.n_qubits))
        )

        if not reduced_str:
            obs = Hermitian(
                matrix=phase * jnp.eye(2, dtype=cdtype()), wires=[0], record=False
            )
            obs._pauli_label = "I"
        else:
            # Reuse the canonical Pauli matrix construction (bare string, then
            # multiply by the leading +-1/+-i phase).
            reduced_word = PauliWord.from_pauli_string(
                reduced_str, list(range(len(reduced_str))), len(reduced_str)
            )
            obs = Hermitian(
                matrix=phase * reduced_word.to_matrix(),
                wires=reduced_wires,
                record=False,
            )
            obs._pauli_label = reduced_str

        obs._pauli_word = word
        return obs

`cliffords_in_observable(operations, original_obs, n_qubits)` `staticmethod` #

Integrates Clifford gates into the observables of the original ansatz, by symbolically conjugating each observable through the final Clifford sequence (O -> C^dagger O C for each Clifford, applied in reverse).

Parameters:

Name	Type	Description	Default
`operations`	`List[Operation]`	Clifford gates	required
`original_obs`	`List[Operation]`	Original observables from the circuit	required
`n_qubits`	`int`	Total number of qubits.	required

Returns:

Type	Description
`List[Operation]`	List[Operation]: Observables with Clifford operations absorbed. Each carries a cached symbolic `_pauli_word` for the Fourier-tree algorithm and a matrix for simulation.

Source code in qml_essentials/pauli.py

@staticmethod
def cliffords_in_observable(
    operations: List[Operation],
    original_obs: List[Operation],
    n_qubits: int,
) -> List[Operation]:
    """
    Integrates Clifford gates into the observables of the original ansatz,
    by symbolically conjugating each observable through the final Clifford
    sequence (``O -> C^dagger O C`` for each Clifford, applied in reverse).

    Args:
        operations (List[Operation]): Clifford gates
        original_obs (List[Operation]): Original observables from the
            circuit
        n_qubits (int): Total number of qubits.

    Returns:
        List[Operation]: Observables with Clifford operations absorbed.
            Each carries a cached symbolic ``_pauli_word`` for the
            Fourier-tree algorithm and a matrix for simulation.
    """
    observables = []
    for ob in original_obs:
        word = PauliWord.from_operation(ob, n_qubits)
        for clifford in operations[::-1]:
            word = word.conjugate_by_clifford(clifford, adjoint_left=True)
        observables.append(PauliCircuit._pauli_operation_from_word(word))
    return observables

`commute_all_cliffords_to_the_end(operations, n_qubits)` `staticmethod` #

This function moves all clifford gates to the end of the circuit, accounting for commutation rules.

Parameters:

Name	Type	Description	Default
`operations`	`List[Operation]`	The operations in the tape of the circuit	required
`n_qubits`	`int`	Total number of qubits.	required

Returns:

Type	Description
`Tuple[List[Operation], List[Operation]]`	Tuple[List[Operation], List[Operation]]: - List of the resulting Pauli-rotations - List of the resulting Clifford gates

Source code in qml_essentials/pauli.py

@staticmethod
def commute_all_cliffords_to_the_end(
    operations: List[Operation],
    n_qubits: int,
) -> Tuple[List[Operation], List[Operation]]:
    """
    This function moves all clifford gates to the end of the circuit,
    accounting for commutation rules.

    Args:
        operations (List[Operation]): The operations in the tape of the
            circuit
        n_qubits (int): Total number of qubits.

    Returns:
        Tuple[List[Operation], List[Operation]]:
            - List of the resulting Pauli-rotations
            - List of the resulting Clifford gates
    """
    first_clifford = -1
    for i in range(len(operations) - 2, -1, -1):
        j = i
        while (
            j + 1 < len(operations)  # Clifford has not alredy reached the end
            and PauliCircuit._is_clifford(operations[j])
            and PauliCircuit._is_pauli_rotation(operations[j + 1])
        ):
            pauli, clifford = PauliCircuit._evolve_clifford_rotation(
                operations[j], operations[j + 1], n_qubits
            )
            operations[j] = pauli
            operations[j + 1] = clifford
            j += 1
            first_clifford = j

    # No Clifford gates are in the circuit
    if not PauliCircuit._is_clifford(operations[-1]):
        return operations, []

    pauli_rotations = operations[:first_clifford]
    clifford_gates = operations[first_clifford:]

    return pauli_rotations, clifford_gates

`from_parameterised_circuit(tape, observables=None, n_qubits=None)` `staticmethod` #

Transforms a list of operations into a Pauli-Clifford circuit.

Parameters:

Name	Type	Description	Default
`tape`	`List[Operation]`	List of operations recorded from the circuit.	required
`observables`	`Optional[List[Operation]]`	List of observable operations. If `None`, defaults to an empty list.	`None`
`n_qubits`	`Optional[int]`	Total number of qubits. Inferred from the maximum wire index if `None`.	`None`

Returns:

Type	Description
`Tuple[List[Operation], List[Operation]]`	Tuple[List[Operation], List[Operation]]: The Pauli rotations of the canonical circuit and the (Clifford-evolved) observables.

Source code in qml_essentials/pauli.py

@staticmethod
def from_parameterised_circuit(
    tape: List[Operation],
    observables: Optional[List[Operation]] = None,
    n_qubits: Optional[int] = None,
) -> Tuple[List[Operation], List[Operation]]:
    """
    Transforms a list of operations into a Pauli-Clifford circuit.

    Args:
        tape: List of operations recorded from the circuit.
        observables: List of observable operations.  If ``None``, defaults
            to an empty list.
        n_qubits: Total number of qubits.  Inferred from the maximum wire
            index if ``None``.

    Returns:
        Tuple[List[Operation], List[Operation]]:
            The Pauli rotations of the canonical circuit and the
            (Clifford-evolved) observables.
    """
    if observables is None:
        observables = []

    operations = PauliCircuit.get_clifford_pauli_gates(tape)

    if n_qubits is None:
        n_qubits = PauliCircuit._infer_n_qubits(operations, observables)

    pauli_gates, final_cliffords = PauliCircuit.commute_all_cliffords_to_the_end(
        operations, n_qubits
    )

    observables = PauliCircuit.cliffords_in_observable(
        final_cliffords, observables, n_qubits
    )

    return pauli_gates, observables

`get_clifford_pauli_gates(tape)` `staticmethod` #

This function decomposes all gates in the circuit to clifford and pauli-rotation gates.

Parameters:

Name	Type	Description	Default
`tape`	`List[Operation]`	List of operations recorded from the circuit.	required

Returns:

Type	Description
`List[Operation]`	List[Operation]: A list of operations consisting only of clifford and Pauli-rotation gates.

Source code in qml_essentials/pauli.py

@staticmethod
def get_clifford_pauli_gates(tape: List[Operation]) -> List[Operation]:
    """
    This function decomposes all gates in the circuit to clifford and
    pauli-rotation gates.

    Args:
        tape: List of operations recorded from the circuit.

    Returns:
        List[Operation]: A list of operations consisting only of clifford
            and Pauli-rotation gates.
    """
    operations = []
    for operation in tape:
        if PauliCircuit._is_clifford(operation) or PauliCircuit._is_pauli_rotation(
            operation
        ):
            operations.append(operation)
        elif PauliCircuit._is_skippable(operation):
            continue
        else:
            # Composite gates (Rot, CRX/CRY/CRZ, ...) expose their own
            # Clifford + Pauli-rotation decomposition.
            try:
                operations.extend(operation.decompose())
            except NotImplementedError:
                raise NotImplementedError(
                    f"Gate {operation.name} cannot be decomposed into "
                    "Pauli rotations and Clifford gates. Consider using a "
                    "circuit ansatz that only uses RX, RY, RZ, PauliRot, "
                    "Rot, and standard Clifford gates."
                )

    return operations

`get_parameters(operations)` `staticmethod` #

Flatten the parameter values of a tape (list of operations).

Source code in qml_essentials/pauli.py

@staticmethod
def get_parameters(operations: List[Operation]) -> list:
    """Flatten the parameter values of a tape (list of operations)."""
    return [p for op in operations for p in op.parameters]

References

Ansaetze#

Matchgate #

Permutation_Equivariant #

XY_Brickwork #

Circuit#

__call__(*args, **kwds) #

__init__() #

build(w, n_qubits, **kwargs) abstractmethod #

get_control_angles(w, n_qubits) #

get_control_indices(n_qubits) abstractmethod #

n_params_per_layer(n_qubits) abstractmethod #

n_pulse_params_per_layer(n_qubits) #

Declarative Circuit#

get_control_indices(n_qubits) classmethod #

structure() classmethod #

Block#

__init__(gate, topology=None, shared=False, wires=None, **kwargs) #

apply(n_qubits, w=None, w_idx=None, **kwargs) #

Encoding#

is_golomb property #

__init__(strategy, gates) #

Parameters#

Returns#

Raises#

binary(enc) #

Parameters#

Returns#

get_n_freqs(data_reupload) #

get_spectrum(data_reupload) #

Parameters#

Returns#

get_weights(n_qubits) #

Parameters#

Returns#

Raises#

golomb(enc) #

Parameters#

Returns#

hamming(enc) #

Parameters#

Returns#

ternary(enc) #

Parameters#

Returns#

Model#

all_qubit_measurement property #

batch_shape property #

data_reupload property writable #

degree property writable #

eff_batch_shape property #

enc_params property writable #

enc_pulse_params property writable #

execution_type property writable #

frequencies property writable #

has_dru property #

noise_params property writable #

observables property writable #

output_qubit property writable #

params property writable #

pulse_params property writable #

shots property writable #

__call__(params=None, inputs=None, pulse_params=None, enc_params=None, data_reupload=None, noise_params=None, execution_type=None, force_mean=False, gate_mode=None, enc_pulse_params=None, random_key=None, keepdims=False) #

__repr__() #

__str__() #

apply(params=None, inputs=None, pulse_params=None, enc_params=None, noise_params=None, execution_type=None, force_mean=False, gate_mode=None, enc_pulse_params=None, key=None) #

draw(inputs=None, figure='text', **kwargs) #

draw_pulse(inputs=None, **kwargs) #

exact_spectrum(method='tree') #

initialize_params(random_key=None, repeat=1, initialization=None, initialization_domain=None) #

next_key() #

transform_input(inputs, enc_params) #

Entanglement#

bell_measurements(model, n_samples, random_key=None, scale=False, **kwargs) classmethod #

concentratable_entanglement(model, n_samples, random_key=None, scale=False, **kwargs) classmethod #

concentratable_entanglement_estimation(model, n_samples, random_key=None, scale=False, **kwargs) classmethod #

entanglement_of_formation(model, n_samples, random_key=None, scale=False, always_decompose=False, **kwargs) classmethod #

meyer_wallach(model, n_samples, random_key=None, scale=False, **kwargs) classmethod #

relative_entropy(model, n_samples, n_sigmas, random_key=None, scale=False, **kwargs) classmethod #

Expressibility#

`Matchgate` #

`Permutation_Equivariant` #

`XY_Brickwork` #

`call(*args, **kwds)` #

`init()` #

`build(w, n_qubits, **kwargs)` `abstractmethod` #

`get_control_angles(w, n_qubits)` #

`get_control_indices(n_qubits)` `abstractmethod` #

`n_params_per_layer(n_qubits)` `abstractmethod` #

`n_pulse_params_per_layer(n_qubits)` #

`get_control_indices(n_qubits)` `classmethod` #

`structure()` `classmethod` #

`init(gate, topology=None, shared=False, wires=None, **kwargs)` #

`apply(n_qubits, w=None, w_idx=None, **kwargs)` #

`is_golomb` `property` #

`init(strategy, gates)` #

`binary(enc)` #

`get_n_freqs(data_reupload)` #

`get_spectrum(data_reupload)` #

`get_weights(n_qubits)` #

`golomb(enc)` #

`hamming(enc)` #

`ternary(enc)` #

`all_qubit_measurement` `property` #

`batch_shape` `property` #

`data_reupload` `property` `writable` #

`degree` `property` `writable` #

`eff_batch_shape` `property` #

`enc_params` `property` `writable` #

`enc_pulse_params` `property` `writable` #

`execution_type` `property` `writable` #

`frequencies` `property` `writable` #

`has_dru` `property` #

`noise_params` `property` `writable` #

`observables` `property` `writable` #

`output_qubit` `property` `writable` #

`params` `property` `writable` #

`pulse_params` `property` `writable` #

`shots` `property` `writable` #

`call(params=None, inputs=None, pulse_params=None, enc_params=None, data_reupload=None, noise_params=None, execution_type=None, force_mean=False, gate_mode=None, enc_pulse_params=None, random_key=None, keepdims=False)` #

`repr()` #

`str()` #

`apply(params=None, inputs=None, pulse_params=None, enc_params=None, noise_params=None, execution_type=None, force_mean=False, gate_mode=None, enc_pulse_params=None, key=None)` #

`draw(inputs=None, figure='text', **kwargs)` #

`draw_pulse(inputs=None, **kwargs)` #

`exact_spectrum(method='tree')` #

`initialize_params(random_key=None, repeat=1, initialization=None, initialization_domain=None)` #

`next_key()` #

`transform_input(inputs, enc_params)` #

`bell_measurements(model, n_samples, random_key=None, scale=False, **kwargs)` `classmethod` #

`concentratable_entanglement(model, n_samples, random_key=None, scale=False, **kwargs)` `classmethod` #

`concentratable_entanglement_estimation(model, n_samples, random_key=None, scale=False, **kwargs)` `classmethod` #

`entanglement_of_formation(model, n_samples, random_key=None, scale=False, always_decompose=False, **kwargs)` `classmethod` #

`meyer_wallach(model, n_samples, random_key=None, scale=False, **kwargs)` `classmethod` #

`relative_entropy(model, n_samples, n_sigmas, random_key=None, scale=False, **kwargs)` `classmethod` #

`haar_integral(n_qubits, n_bins, cache=True, scale=False)` `classmethod` #

`kl_divergence_to_haar(model, n_samples, n_bins, random_key=None, scale=False, **kwargs)` `classmethod` #

`kullback_leibler_divergence(vqc_prob_dist, haar_dist)` `classmethod` #

`state_fidelities(n_samples, n_bins, model, random_key=None, scale=False, **kwargs)` `classmethod` #

`evaluate_Fourier_series(coefficients, frequencies, inputs)` `classmethod` #

`get_psd(coeffs)` `classmethod` #

`get_spectrum(model, mfs=1, mts=1, shift=False, trim=False, numerical_cap=-1, **kwargs)` `classmethod` #

`call(params=None, inputs=None, **kwargs)` #

`init(model)` #

`get_exact_support(method='tree')` #

`get_spectrum(force_mean=False)` #

`calculate_fcc(fourier_fingerprint)` `classmethod` #

`get_fcc(model, n_samples, random_key=None, method='pearson', scale=False, weight=False, trim_redundant=True, **kwargs)` `classmethod` #

`get_fourier_fingerprint(model, n_samples, random_key=None, method='pearson', scale=False, weight=False, trim_redundant=True, nan_to_one=False, **kwargs)` `classmethod` #

`calculate_values(domain_samples, frequencies, coefficients)` `classmethod` #

`construct_coefficients(random_key, model, coefficients_min=0.0, coefficients_max=1.0, zero_centered=False)` `classmethod` #

`construct_domain_samples(model, mts=1, mfs=1)` `classmethod` #

`construct_frequencies(model, random_key=None, offgrid_mode='none', offgrid_prob=0.0, offgrid_resolution=2)` `classmethod` #

`generate_fourier_series(random_key, model, coefficients_min=0.0, coefficients_max=1.0, zero_centered=False)` `classmethod` #

`generator_etas(model, random_key, offgrid_prob, offgrid_resolution)` `classmethod` #

`uniform_circle(random_key, size, low=0.0, high=1.0)` `classmethod` #

`all_pairs(n_qubits)` `classmethod` #

`all_to_all(n_qubits)` `classmethod` #

`graph(n_qubits, *, edges)` `classmethod` #

`stairs(n_qubits, offset=0, wrap=False, reverse=True, mirror=True, span=1, stride=1, modulo=True)` `classmethod` #